Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
635679c13d | ||
|
|
682b9b8044 | ||
|
|
8b26776cdd | ||
|
|
5c6a4a40ec | ||
|
|
73ff7b6bb2 | ||
|
|
7f4f0d50bf | ||
|
|
5a59f7ae22 | ||
|
|
a452aa16a1 | ||
|
|
be48ecb550 | ||
|
|
484c33f986 | ||
|
|
692a334c09 | ||
|
|
72e0828360 | ||
|
|
0d923bc366 | ||
|
|
e0168f7d65 | ||
|
|
816a76fe5b | ||
|
|
0e691a6b4f | ||
|
|
38f0689ae3 | ||
|
|
e2ddb2e9d7 | ||
|
|
013a39cb4b | ||
|
|
103ec3d05a | ||
|
|
2529cb8e84 | ||
|
|
583a35883d | ||
|
|
94da9d9127 | ||
|
|
bdc093ee02 | ||
|
|
6e0949d58f | ||
|
|
a804f5ac93 | ||
|
|
3a3a470cd7 | ||
|
|
7c40531891 | ||
|
|
0481b6e32f | ||
|
|
64d08a358d | ||
|
|
8cdb3e8537 | ||
|
|
f6644c9278 | ||
|
|
300b46bb67 | ||
|
|
7b7eb4e19b | ||
|
|
7595ad8519 | ||
|
|
d52a4fa1e6 | ||
|
|
218eaa8b53 | ||
|
|
63e70e19fc | ||
|
|
b1da18d49e | ||
|
|
49d96e197b | ||
|
|
1311cf38d4 | ||
|
|
c89281200a | ||
|
|
01c8d9daeb |
@@ -18,6 +18,10 @@ Go and runs on orion under a hardened systemd unit. You are implementing it one
|
||||
only that file, and tell the owner.
|
||||
5. Work only from files inside this repository. Never read another checkout, and never search the
|
||||
file system for code. If you are stuck, stop as in point 4.
|
||||
6. **Write code into files early and in small steps.** Start each new file with its types and
|
||||
signatures, make it compile, then fill in one function at a time and run the tests as you go.
|
||||
Never write a whole implementation out in your reply before putting it in a file: a long
|
||||
session is compacted, and what was only in the reply is lost (P6 task 02, 2026-09-26).
|
||||
|
||||
## Files you must never edit
|
||||
|
||||
|
||||
@@ -4,7 +4,93 @@ LLM-oriented fetch and search for the Hermes fleet: fetch a resource and get it
|
||||
compact, and with useful context attached — using local models (Dixie), SearXNG and Brave on
|
||||
orion, and Honcho for retrieval.
|
||||
|
||||
Status: P1 clean fetch (`/v2/scrape`, `/v1/fetch`, CLI, auth). See [`docs/PLAN.md`](docs/PLAN.md) and [`docs/P1.md`](docs/P1.md).
|
||||
Status: implemented through P5a: P1 clean fetch (`/v2/scrape`, `/v1/fetch`, CLI, auth), P2
|
||||
digest, fleet context and recover, P3 search, P2.7 Prometheus metrics, P5a PDF documents. The design
|
||||
is [`docs/PLAN.md`](docs/PLAN.md); each phase has its own doc (`docs/P1.md` … `docs/P5a.md`).
|
||||
|
||||
## API
|
||||
|
||||
All endpoints except `/healthz` take `Authorization: Bearer <key>` (per caller; see `sift keys`).
|
||||
The server binds `127.0.0.1:8890` and is published on the tailnet by `tailscale serve` at
|
||||
`https://orion.scylla-hammerhead.ts.net:8443`.
|
||||
|
||||
| Endpoint | What |
|
||||
|---|---|
|
||||
| `POST /v2/scrape` | Firecrawl-compatible fetch (`url`, `formats`) — what Hermes `web_extract` calls |
|
||||
| `POST /v2/search` | Firecrawl-compatible search (`query`, `limit`) — what Hermes `web_search` calls |
|
||||
| `POST /v1/fetch` | native fetch: `url`, `mode` (`clean`/`raw`/`digest`), `question`, `jq`, `context`, `recover`, and for PDFs `pages`, `grep` |
|
||||
| `POST /v1/search` | native search: `query` or `queries`, `k`, `rerank`, `digest`, domain/time/topic/`exact_match`/`min_score` |
|
||||
| `GET /metrics`, `/metrics.json` | Prometheus text or JSON; open on loopback, bearer elsewhere |
|
||||
| `GET /healthz` | liveness |
|
||||
|
||||
Every outbound connection goes through the SSRF guard (`internal/guard`): private, tailnet and
|
||||
configured `deny_prefixes` addresses, and the host's own WAN address (`[self_ip]`), are refused
|
||||
with `blocked: …`.
|
||||
|
||||
## CLI
|
||||
|
||||
```
|
||||
sift serve [-listen addr] [-config /etc/sift/config.toml] [-keys /var/lib/sift/keys.toml] [-deny-prefix CIDR]...
|
||||
sift fetch <url> [--question q] [--mode clean|raw|digest] [--jq filter] [--pages a-b] [--grep re]
|
||||
[--no-context] [--recover] [--json]
|
||||
sift extract <file.pdf> [--pages a-b] [--grep re] [--question q] [--mode clean|digest] [--json]
|
||||
sift search "<query>" [--k n] [--json] [--include-domain host]... [--topic news] [--time-range month]
|
||||
sift batch submit|status|get|cancel
|
||||
sift keys add --caller name [--workspace ws] [--out file] | list | revoke --caller name
|
||||
sift version
|
||||
```
|
||||
|
||||
`sift fetch`, `sift extract` and `sift search` are clients: they read `url` and `key_file` from
|
||||
`client.toml` (`$SIFT_CLIENT`, else `$HERMES_HOME/sift/client.toml`, else `~/.config/sift/client.toml`);
|
||||
the key file must be mode 0600. `sift extract` uploads a local PDF to `/v1/extract` and prints the
|
||||
answer exactly as `sift fetch` does for a PDF URL — `--pages`, `--grep` and `--question` behave the
|
||||
same, and a file over 100 MiB is refused before the upload. The agent-facing guide is
|
||||
[`skill/SKILL.md`](skill/SKILL.md). JSON `search_queries` (and `sift_search_queries_*`) counts
|
||||
each sub-query as its own upstream call, so a 4-query search is 4× one `query`.
|
||||
|
||||
## Batch
|
||||
|
||||
`sift batch` runs several fetches as one job, so an agent can read a handful of URLs later instead
|
||||
of a loop of `sift fetch`. The four endpoints:
|
||||
|
||||
| Endpoint | What |
|
||||
|---|---|
|
||||
| `POST /v1/batch` | `{"targets":[…]}` → `202 {"success","job_id","status","targets","reused"}` |
|
||||
| `GET /v1/batch/{id}` | summary; `?results=1` inlines each done target's `/v1/fetch` envelope |
|
||||
| `GET /v1/batch/{id}/{idx}` | one target's result, as `sift fetch --json` |
|
||||
| `DELETE /v1/batch/{id}` | cancel a queued job → `200 {"success","status"}` |
|
||||
|
||||
`sift batch submit <targets.json|->` reads a JSON object with a non-empty `targets` array locally,
|
||||
or a bare array of targets or URL strings (each string becomes `{"url": s}`), then posts it. `status`
|
||||
prints the summary (or the body with `?results=1`); `get <job> <idx>` prints one result; `cancel`
|
||||
stops a queued job.
|
||||
|
||||
A job runs **serially** (target `i+1` never starts before `i` is stored) with **one worker by
|
||||
default** (see `[batch]` config) and the runner rotates jobs **round-robin** while several are live.
|
||||
Caps: **32 targets** per job, **1 live job per caller**, **8 live jobs** in all. The posted body is
|
||||
capped at **256 KiB**. A target that is blocked, times out, or errors is recorded and the job
|
||||
continues; a job whose targets are all done is `done`, all failed `failed`, mixed `partial`.
|
||||
|
||||
**Idempotency**: the same list from the same caller returns the same job while it is live or already
|
||||
has results (`reused: true`), so a retry does not start a second job. Each job carries a **lease**
|
||||
that a status read renews; a job whose lease has expired for more than the **24 h TTL** is swept
|
||||
along with its targets. On restart, jobs in flight when the server died are marked failed so a retry
|
||||
re-submits them.
|
||||
|
||||
A `status` read is a **summary** (per-target status, never the content); `?results=1` inlines the
|
||||
full envelope for each done target, so it can be large — prefer `get <job> <idx>` for one result.
|
||||
Restricted callers see the same endpoints; their results are **redacted** the same way a restricted
|
||||
`sift fetch` is. Batch counters are `sift_batch_jobs_total{outcome}`,
|
||||
`sift_batch_targets_total{outcome}`, `sift_batch_queue_depth`, `sift_batch_targets_inflight`,
|
||||
`sift_batch_job_seconds`, `sift_batch_target_seconds` and `sift_batch_result_bytes`.
|
||||
|
||||
## Configuration
|
||||
|
||||
`/etc/sift/config.toml` (`$SIFT_CONFIG`); every key is shown in
|
||||
[`dist/sift.toml.example`](dist/sift.toml.example). Paid-search keys (`BRAVE_API_KEY`,
|
||||
`TAVILY_API_KEY`) live on the box in `/var/lib/sift/secret.env` (`$SIFT_SECRET`), never in git or
|
||||
the Nix store. On NixOS use `nixosModules.sift` (`services.sift`), which writes the config and
|
||||
the hardened unit.
|
||||
|
||||
## Build
|
||||
|
||||
@@ -42,6 +128,13 @@ Answers (through `/v1/fetch` or `/v2/scrape`):
|
||||
- `--pages 41-60` for a range, `--grep 'pattern'` for matching lines with page context;
|
||||
- `--question "…"` for a **digest** that, on a PDF, cites each fact as `(p. N)`.
|
||||
|
||||
A range returns at most 100 pages; grep shows at most 100 matches and takes a pattern of at most
|
||||
200 bytes. `pages` and `grep` on anything but a PDF are a `400`. A question on a long document is
|
||||
answered first from BM25 (`retrieval: "bm25"`) — lexical, before the document's chunks are
|
||||
embedded in the background — and hybrid on later questions, once those embeddings exist. A question
|
||||
over a document too long to answer even by retrieval (more than 5000 chunks) is refused with
|
||||
`422 document too long for question (N chunks); use grep or pages`.
|
||||
|
||||
Scanned PDFs have no text layer and return `no text layer (scanned PDF?)`; point them at an OCR
|
||||
source rather than fixing them here.
|
||||
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
package main
|
||||
|
||||
// P6 owner review: the batch CLI treats every 2xx as success. POST /v1/batch answers 202, and
|
||||
// batchDo exited 1 on it with an empty message.
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBatchDoAccepts202(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Header.Get("Authorization") != "Bearer k" {
|
||||
w.WriteHeader(401)
|
||||
_, _ = io.WriteString(w, `{"error":"unauthorized"}`)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusAccepted)
|
||||
_, _ = io.WriteString(w, `{"success":true,"job_id":"j1","status":"queued","targets":1,"reused":false}`)
|
||||
}))
|
||||
defer srv.Close()
|
||||
dir := t.TempDir()
|
||||
keyPath := filepath.Join(dir, "key")
|
||||
if err := os.WriteFile(keyPath, []byte("k\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg := filepath.Join(dir, "client.toml")
|
||||
if err := os.WriteFile(cfg, []byte("url = \""+srv.URL+"\"\nkey_file = \""+keyPath+"\"\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Setenv("SIFT_CLIENT", cfg)
|
||||
|
||||
r, w, err := os.Pipe()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
old := os.Stdout
|
||||
os.Stdout = w
|
||||
batchDo(http.MethodPost, "/v1/batch", strings.NewReader(`{"targets":[{"url":"https://example.com/"}]}`))
|
||||
os.Stdout = old
|
||||
_ = w.Close()
|
||||
out, _ := io.ReadAll(r)
|
||||
if !strings.Contains(string(out), `"job_id": "j1"`) {
|
||||
t.Fatalf("stdout %q, want the 202 body printed", out)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestParseBatchArgs(t *testing.T) {
|
||||
cases := []struct {
|
||||
args []string
|
||||
verb string
|
||||
rest []string
|
||||
ok bool
|
||||
}{
|
||||
{[]string{"submit", "targets.json"}, "submit", []string{"targets.json"}, true},
|
||||
{[]string{"submit", "-"}, "submit", []string{"-"}, true},
|
||||
{[]string{"status", "job-123"}, "status", []string{"job-123"}, true},
|
||||
{[]string{"status", "--results", "job-123"}, "status", []string{"--results", "job-123"}, true},
|
||||
{[]string{"get", "job-123", "4"}, "get", []string{"job-123", "4"}, true},
|
||||
{[]string{"cancel", "job-123"}, "cancel", []string{"job-123"}, true},
|
||||
{[]string{}, "", nil, false},
|
||||
{[]string{"submit"}, "", nil, false},
|
||||
{[]string{"status"}, "", nil, false},
|
||||
{[]string{"status", "a", "b"}, "", nil, false},
|
||||
{[]string{"status", "--results"}, "", nil, false},
|
||||
{[]string{"get", "job-123"}, "", nil, false},
|
||||
{[]string{"get", "job-123", "x"}, "", nil, false},
|
||||
{[]string{"get", "job-123", "-1"}, "", nil, false},
|
||||
{[]string{"cancel"}, "", nil, false},
|
||||
{[]string{"bogus", "x"}, "", nil, false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
verb, rest, err := parseBatchArgs(c.args)
|
||||
if c.ok {
|
||||
if err != nil {
|
||||
t.Errorf("%v: unexpected error %v", c.args, err)
|
||||
continue
|
||||
}
|
||||
if verb != c.verb || !slices.Equal(rest, c.rest) {
|
||||
t.Errorf("%v: got verb=%q rest=%q", c.args, verb, rest)
|
||||
}
|
||||
} else if err == nil {
|
||||
t.Errorf("%v: want an error", c.args)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,200 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
const batchUsage = `usage: sift batch submit|status|get|cancel
|
||||
sift batch submit <targets.json|->
|
||||
sift batch status [--results] <job_id>
|
||||
sift batch get <job_id> <idx>
|
||||
sift batch cancel <job_id>`
|
||||
|
||||
// parseBatchArgs validates a `sift batch` command line (without "batch"). verb is
|
||||
// the subcommand; rest is the arguments after it, unchanged.
|
||||
func parseBatchArgs(args []string) (verb string, rest []string, err error) {
|
||||
if len(args) == 0 {
|
||||
return "", nil, fmt.Errorf("%s", batchUsage)
|
||||
}
|
||||
verb, rest = args[0], args[1:]
|
||||
switch verb {
|
||||
case "submit":
|
||||
if len(rest) != 1 {
|
||||
return "", nil, fmt.Errorf("%s", batchUsage)
|
||||
}
|
||||
case "status":
|
||||
switch len(rest) {
|
||||
case 1:
|
||||
if rest[0] == "--results" {
|
||||
return "", nil, fmt.Errorf("%s", batchUsage)
|
||||
}
|
||||
case 2:
|
||||
if rest[0] != "--results" {
|
||||
return "", nil, fmt.Errorf("%s", batchUsage)
|
||||
}
|
||||
default:
|
||||
return "", nil, fmt.Errorf("%s", batchUsage)
|
||||
}
|
||||
case "get":
|
||||
if len(rest) != 2 {
|
||||
return "", nil, fmt.Errorf("%s", batchUsage)
|
||||
}
|
||||
idx, e := strconv.Atoi(rest[1])
|
||||
if e != nil || idx < 0 {
|
||||
return "", nil, fmt.Errorf("%s", batchUsage)
|
||||
}
|
||||
case "cancel":
|
||||
if len(rest) != 1 {
|
||||
return "", nil, fmt.Errorf("%s", batchUsage)
|
||||
}
|
||||
default:
|
||||
return "", nil, fmt.Errorf("%s", batchUsage)
|
||||
}
|
||||
return verb, rest, nil
|
||||
}
|
||||
|
||||
func batchMain() {
|
||||
verb, rest, err := parseBatchArgs(os.Args[1:])
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(2)
|
||||
}
|
||||
switch verb {
|
||||
case "submit":
|
||||
submitMain(rest)
|
||||
case "status":
|
||||
statusMain(rest)
|
||||
case "get":
|
||||
getMain(rest)
|
||||
case "cancel":
|
||||
cancelMain(rest)
|
||||
}
|
||||
}
|
||||
|
||||
// submitMain reads the targets file (or stdin when the argument is "-") and
|
||||
// validates it locally before calling the server: a JSON object with a
|
||||
// non-empty targets array, or a bare array of targets or URL strings wrapped
|
||||
// into {"targets": [...]}.
|
||||
func submitMain(rest []string) {
|
||||
var r io.Reader = os.Stdin
|
||||
if rest[0] != "-" {
|
||||
f, err := os.Open(rest[0])
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(5)
|
||||
}
|
||||
defer f.Close()
|
||||
r = f
|
||||
}
|
||||
raw, err := io.ReadAll(r)
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(5)
|
||||
}
|
||||
body, err := buildTargetsBody(raw)
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(2)
|
||||
}
|
||||
batchDo(http.MethodPost, "/v1/batch", bytes.NewReader(body))
|
||||
}
|
||||
|
||||
// buildTargetsBody accepts the three accepted shapes and returns the
|
||||
// {"targets": [...]} body to send. A bare array of URL strings gets each string
|
||||
// wrapped as {"url": s}; bare target objects pass through.
|
||||
func buildTargetsBody(raw []byte) ([]byte, error) {
|
||||
var probe any
|
||||
if err := json.Unmarshal(raw, &probe); err != nil {
|
||||
return nil, fmt.Errorf("targets: not valid JSON: %w", err)
|
||||
}
|
||||
switch p := probe.(type) {
|
||||
case map[string]any:
|
||||
arr, ok := p["targets"].([]any)
|
||||
if !ok || len(arr) == 0 {
|
||||
return nil, fmt.Errorf("targets: a JSON object needs a non-empty targets array")
|
||||
}
|
||||
return raw, nil
|
||||
case []any:
|
||||
if len(p) == 0 {
|
||||
return nil, fmt.Errorf("targets: an empty array")
|
||||
}
|
||||
out := make([]any, 0, len(p))
|
||||
for _, el := range p {
|
||||
switch e := el.(type) {
|
||||
case map[string]any:
|
||||
out = append(out, e)
|
||||
case string:
|
||||
out = append(out, map[string]any{"url": e})
|
||||
default:
|
||||
return nil, fmt.Errorf("targets: array elements must be URLs or objects")
|
||||
}
|
||||
}
|
||||
return json.Marshal(map[string]any{"targets": out})
|
||||
default:
|
||||
return nil, fmt.Errorf("targets: expected a JSON object or array")
|
||||
}
|
||||
}
|
||||
|
||||
func statusMain(rest []string) {
|
||||
id := rest[len(rest)-1]
|
||||
path := "/v1/batch/" + url.PathEscape(id)
|
||||
if len(rest) == 2 && rest[0] == "--results" {
|
||||
path += "?results=1"
|
||||
}
|
||||
batchDo(http.MethodGet, path, nil)
|
||||
}
|
||||
|
||||
func getMain(rest []string) {
|
||||
path := "/v1/batch/" + url.PathEscape(rest[0]) + "/" + url.PathEscape(rest[1])
|
||||
batchDo(http.MethodGet, path, nil)
|
||||
}
|
||||
|
||||
func cancelMain(rest []string) {
|
||||
batchDo(http.MethodDelete, "/v1/batch/"+url.PathEscape(rest[0]), nil)
|
||||
}
|
||||
|
||||
// batchDo calls base+path with the bearer header, decodes the JSON envelope,
|
||||
// and prints it on stdout. On a non-2xx it writes the error to stderr and exits
|
||||
// with exitFor, as fetch does.
|
||||
func batchDo(method, path string, body io.Reader) {
|
||||
base, key, err := loadClient()
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(5)
|
||||
}
|
||||
req, err := http.NewRequest(method, base+path, body)
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(5)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
if body != nil {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(5)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
var env map[string]any
|
||||
if err := json.NewDecoder(resp.Body).Decode(&env); err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(5)
|
||||
}
|
||||
if resp.StatusCode < 200 || resp.StatusCode > 299 { // submit answers 202 Accepted
|
||||
msg, _ := env["error"].(string)
|
||||
fmt.Fprintln(os.Stderr, msg)
|
||||
os.Exit(exitFor(resp.StatusCode, msg))
|
||||
}
|
||||
enc := json.NewEncoder(os.Stdout)
|
||||
enc.SetIndent("", " ")
|
||||
_ = enc.Encode(env)
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
package main
|
||||
|
||||
import "testing"
|
||||
|
||||
// P5c owner review: a client base URL with a path prefix keeps it, as fetch and search do.
|
||||
func TestExtractURLKeepsBasePath(t *testing.T) {
|
||||
got := extractURL("https://sift.example:8443/prefix/", extractOpts{File: "a.pdf", Pages: "2"})
|
||||
if want := "https://sift.example:8443/prefix/v1/extract?pages=2"; got != want {
|
||||
t.Fatalf("extractURL = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
package main
|
||||
|
||||
// P5c task 02: `sift extract <file>` uploads a PDF to /v1/extract.
|
||||
|
||||
import (
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestParseExtractOpts(t *testing.T) {
|
||||
o, err := parseExtractOpts([]string{"thesis.pdf", "--pages", "41-60", "--grep", "attention head"})
|
||||
if err != nil || o.File != "thesis.pdf" || o.Pages != "41-60" || o.Grep != "attention head" {
|
||||
t.Errorf("%+v %v", o, err)
|
||||
}
|
||||
o, err = parseExtractOpts([]string{"--question", "what is the method?", "--json", "paper.pdf"})
|
||||
if err != nil || o.File != "paper.pdf" || o.Question != "what is the method?" || !o.JSON {
|
||||
t.Errorf("flags first: %+v %v", o, err)
|
||||
}
|
||||
if _, err := parseExtractOpts(nil); err == nil {
|
||||
t.Error("no file: want an error")
|
||||
}
|
||||
if _, err := parseExtractOpts([]string{"a.pdf", "b.pdf"}); err == nil {
|
||||
t.Error("two files: want an error")
|
||||
}
|
||||
if !strings.Contains(extractUsage, "sift extract <file.pdf>") {
|
||||
t.Errorf("usage %q", extractUsage)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractURL(t *testing.T) {
|
||||
u, err := url.Parse(extractURL("https://sift.example:8443", extractOpts{File: "x.pdf", Pages: "2-3", Grep: "a b&c"}))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if u.Path != "/v1/extract" || u.Query().Get("pages") != "2-3" || u.Query().Get("grep") != "a b&c" {
|
||||
t.Errorf("url %s", u)
|
||||
}
|
||||
if u.Query().Has("question") || u.Query().Has("mode") {
|
||||
t.Errorf("unset options must not be sent: %s", u)
|
||||
}
|
||||
u, _ = url.Parse(extractURL("https://sift.example:8443", extractOpts{File: "x.pdf", Question: "why?"}))
|
||||
if u.Query().Get("question") != "why?" || u.Query().Get("mode") != "digest" {
|
||||
t.Errorf("question url %s (mode must be digest with a question)", u)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,156 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// maxExtractUpload mirrors the server's /v1/extract cap (a file over 100 MiB is
|
||||
// refused on upload with 413), so reject it before uploading.
|
||||
const maxExtractUpload = 100 << 20
|
||||
|
||||
type extractOpts struct {
|
||||
File, Pages, Grep, Question, Mode string
|
||||
JSON bool
|
||||
}
|
||||
|
||||
const extractUsage = "usage: sift extract <file.pdf> [--pages a-b] [--grep RE] [--question Q] [--mode clean|digest] [--json]"
|
||||
|
||||
// parseExtractOpts parses `sift extract` args into extractOpts. It knows which
|
||||
// flags take a value; the PDF path is the one positional.
|
||||
func parseExtractOpts(args []string) (extractOpts, error) {
|
||||
flagArgs, pos, err := splitAroundPositional(args, extractValueFlags)
|
||||
if err != nil {
|
||||
return extractOpts{}, err
|
||||
}
|
||||
fs := flag.NewFlagSet("extract", flag.ContinueOnError)
|
||||
fs.SetOutput(io.Discard)
|
||||
modeF := fs.String("mode", "", "clean|digest (default digest if --question)")
|
||||
pagesF := fs.String("pages", "", "only these pages of a PDF (a-b or a)")
|
||||
grepF := fs.String("grep", "", "only lines matching RE in a PDF")
|
||||
questionF := fs.String("question", "", "focus the digest")
|
||||
asJSONF := fs.Bool("json", false, "print full /v1/extract envelope")
|
||||
if err := fs.Parse(flagArgs); err != nil {
|
||||
return extractOpts{}, err
|
||||
}
|
||||
if fs.NArg() != 0 {
|
||||
return extractOpts{}, fmt.Errorf("%s", extractUsage)
|
||||
}
|
||||
if len(pos) != 1 {
|
||||
return extractOpts{}, fmt.Errorf("%s", extractUsage)
|
||||
}
|
||||
return extractOpts{
|
||||
File: pos[0],
|
||||
Mode: *modeF,
|
||||
Pages: *pagesF,
|
||||
Grep: *grepF,
|
||||
Question: *questionF,
|
||||
JSON: *asJSONF,
|
||||
}, nil
|
||||
}
|
||||
|
||||
var extractValueFlags = map[string]bool{
|
||||
"mode": true, "pages": true, "grep": true, "question": true,
|
||||
}
|
||||
|
||||
// extractURL builds the /v1/extract URL from base plus only the options that are
|
||||
// set; mode defaults to digest when a question is given and no mode was.
|
||||
func extractURL(base string, o extractOpts) string {
|
||||
v := url.Values{}
|
||||
if o.Pages != "" {
|
||||
v.Set("pages", o.Pages)
|
||||
}
|
||||
if o.Grep != "" {
|
||||
v.Set("grep", o.Grep)
|
||||
}
|
||||
if o.Question != "" {
|
||||
v.Set("question", o.Question)
|
||||
}
|
||||
mode := o.Mode
|
||||
if mode == "" && o.Question != "" {
|
||||
mode = "digest"
|
||||
}
|
||||
if mode != "" {
|
||||
v.Set("mode", mode)
|
||||
}
|
||||
u := strings.TrimRight(base, "/") + "/v1/extract"
|
||||
if q := v.Encode(); q != "" {
|
||||
u += "?" + q
|
||||
}
|
||||
return u
|
||||
}
|
||||
|
||||
func extractMain() {
|
||||
o, err := parseExtractOpts(os.Args[1:])
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(2)
|
||||
}
|
||||
base, key, err := loadClient()
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(5)
|
||||
}
|
||||
f, err := os.Open(o.File)
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(5)
|
||||
}
|
||||
defer f.Close()
|
||||
st, err := f.Stat()
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(5)
|
||||
}
|
||||
if st.Size() > maxExtractUpload {
|
||||
fmt.Fprintf(os.Stderr, "extract: %s is %d bytes, over the %d MB limit\n", o.File, st.Size(), maxExtractUpload>>20)
|
||||
os.Exit(3)
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodPost, extractURL(base, o), f)
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(5)
|
||||
}
|
||||
req.ContentLength = st.Size()
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
req.Header.Set("Content-Type", "application/pdf")
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(5)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
var env map[string]any
|
||||
if err := json.NewDecoder(resp.Body).Decode(&env); err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(5)
|
||||
}
|
||||
if resp.StatusCode != 200 {
|
||||
msg, _ := env["error"].(string)
|
||||
hint, _ := env["hint"].(string)
|
||||
if hint != "" {
|
||||
fmt.Fprintf(os.Stderr, "%s\nhint: %s\n", msg, hint)
|
||||
} else {
|
||||
fmt.Fprintln(os.Stderr, msg)
|
||||
}
|
||||
os.Exit(exitFor(resp.StatusCode, msg))
|
||||
}
|
||||
if o.JSON {
|
||||
enc := json.NewEncoder(os.Stdout)
|
||||
enc.SetIndent("", " ")
|
||||
_ = enc.Encode(env)
|
||||
return
|
||||
}
|
||||
if c, ok := env["content"].(string); ok {
|
||||
fmt.Print(c)
|
||||
if !strings.HasSuffix(c, "\n") {
|
||||
fmt.Println()
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -22,8 +22,12 @@ func main() {
|
||||
serveMain()
|
||||
case "fetch":
|
||||
fetchMain()
|
||||
case "extract":
|
||||
extractMain()
|
||||
case "search":
|
||||
searchMain()
|
||||
case "batch":
|
||||
batchMain()
|
||||
case "keys":
|
||||
keysMain()
|
||||
case "pdftext":
|
||||
@@ -45,7 +49,9 @@ func usage() {
|
||||
Commands:
|
||||
sift serve [flags] run the service
|
||||
sift fetch <url> fetch via the service
|
||||
sift extract <file.pdf> extract a local PDF via the service
|
||||
sift search "<query>" search via the service
|
||||
sift batch submit|status|get|cancel run or watch a batch of fetches
|
||||
sift keys add|list|revoke
|
||||
sift version
|
||||
`, version)
|
||||
|
||||
+89
-7
@@ -8,15 +8,35 @@ import (
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type stringList []string
|
||||
|
||||
func (s *stringList) String() string { return strings.Join(*s, ",") }
|
||||
func (s *stringList) Set(v string) error {
|
||||
*s = append(*s, v)
|
||||
return nil
|
||||
}
|
||||
|
||||
func searchMain() {
|
||||
fs := flag.NewFlagSet("search", flag.ContinueOnError)
|
||||
fs.SetOutput(io.Discard)
|
||||
k := fs.Int("k", 0, "result cap")
|
||||
asJSON := fs.Bool("json", false, "native envelope")
|
||||
var include, exclude stringList
|
||||
fs.Var(&include, "include-domain", "restrict or prefer this host")
|
||||
fs.Var(&exclude, "exclude-domain", "drop this host")
|
||||
domainsMode := fs.String("domains-mode", "", "restrict|prefer")
|
||||
timeRange := fs.String("time-range", "", "day|week|month|year")
|
||||
startDate := fs.String("start-date", "", "YYYY-MM-DD")
|
||||
endDate := fs.String("end-date", "", "YYYY-MM-DD")
|
||||
topic := fs.String("topic", "", "general|news|finance")
|
||||
exact := fs.Bool("exact-match", false, "require quoted phrases")
|
||||
minScore := fs.Float64("min-score", 0, "Tavily score cut")
|
||||
queriesFile := fs.String("queries-file", "", "JSON array of 1-4 queries")
|
||||
args := os.Args[1:]
|
||||
flagArgs, pos, err := splitAroundPositional(args, map[string]bool{"k": true})
|
||||
flagArgs, pos, err := splitAroundPositional(args, searchValueFlags)
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(2)
|
||||
@@ -25,20 +45,65 @@ func searchMain() {
|
||||
fmt.Fprintln(os.Stderr, searchUsage)
|
||||
os.Exit(2)
|
||||
}
|
||||
body := map[string]any{}
|
||||
if *queriesFile != "" {
|
||||
if len(pos) != 0 {
|
||||
fmt.Fprintln(os.Stderr, searchUsage)
|
||||
os.Exit(2)
|
||||
}
|
||||
raw, err := os.ReadFile(*queriesFile)
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(2)
|
||||
}
|
||||
var qs []string
|
||||
if err := json.Unmarshal(raw, &qs); err != nil {
|
||||
fmt.Fprintln(os.Stderr, "queries-file: want JSON array of strings")
|
||||
os.Exit(2)
|
||||
}
|
||||
body["queries"] = qs
|
||||
} else {
|
||||
if len(pos) != 1 {
|
||||
fmt.Fprintln(os.Stderr, searchUsage)
|
||||
os.Exit(2)
|
||||
}
|
||||
query := pos[0]
|
||||
body["query"] = pos[0]
|
||||
}
|
||||
if *k > 0 {
|
||||
body["k"] = *k
|
||||
}
|
||||
if len(include) > 0 {
|
||||
body["include_domains"] = []string(include)
|
||||
}
|
||||
if len(exclude) > 0 {
|
||||
body["exclude_domains"] = []string(exclude)
|
||||
}
|
||||
if *domainsMode != "" {
|
||||
body["include_domains_mode"] = *domainsMode
|
||||
}
|
||||
if *timeRange != "" {
|
||||
body["time_range"] = *timeRange
|
||||
}
|
||||
if *startDate != "" {
|
||||
body["start_date"] = *startDate
|
||||
}
|
||||
if *endDate != "" {
|
||||
body["end_date"] = *endDate
|
||||
}
|
||||
if *topic != "" {
|
||||
body["topic"] = *topic
|
||||
}
|
||||
if *exact {
|
||||
body["exact_match"] = true
|
||||
}
|
||||
if *minScore != 0 {
|
||||
body["min_score"] = *minScore
|
||||
}
|
||||
base, key, err := loadClient()
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(5)
|
||||
}
|
||||
body := map[string]any{"query": query}
|
||||
if *k > 0 {
|
||||
body["k"] = *k
|
||||
}
|
||||
raw, _ := json.Marshal(body)
|
||||
req, err := http.NewRequest(http.MethodPost, base+"/v1/search", bytes.NewReader(raw))
|
||||
if err != nil {
|
||||
@@ -80,7 +145,24 @@ func searchMain() {
|
||||
if sn, _ := m["snippet"].(string); sn != "" {
|
||||
fmt.Println(sn)
|
||||
}
|
||||
if d, _ := m["published_date"].(string); d != "" {
|
||||
fmt.Println(d)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const searchUsage = `usage: sift search "<query>" [--k N] [--json]`
|
||||
var searchValueFlags = map[string]bool{
|
||||
"k": true, "include-domain": true, "exclude-domain": true,
|
||||
"domains-mode": true, "time-range": true, "start-date": true, "end-date": true,
|
||||
"topic": true, "min-score": true, "queries-file": true,
|
||||
}
|
||||
|
||||
const searchUsage = `usage: sift search "<query>" [--k n] [--json]
|
||||
[--include-domain host]... [--exclude-domain host]...
|
||||
[--domains-mode restrict|prefer]
|
||||
[--time-range day|week|month|year]
|
||||
[--start-date YYYY-MM-DD] [--end-date YYYY-MM-DD]
|
||||
[--topic general|news|finance]
|
||||
[--exact-match]
|
||||
[--min-score 0.5]
|
||||
[--queries-file path]`
|
||||
|
||||
+26
-1
@@ -10,9 +10,11 @@ import (
|
||||
"os/signal"
|
||||
"strings"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/api"
|
||||
"git.wntrmute.dev/kyle/sift/internal/auth"
|
||||
"git.wntrmute.dev/kyle/sift/internal/batch"
|
||||
"git.wntrmute.dev/kyle/sift/internal/config"
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
"git.wntrmute.dev/kyle/sift/internal/digest"
|
||||
@@ -119,7 +121,7 @@ func serveMain() {
|
||||
cachePath = "/var/lib/sift/cache.db"
|
||||
}
|
||||
var cache *dcache.DB
|
||||
if cfg.Dixie.Chat != "" || cfg.Search.AnyEnabled() {
|
||||
if cfg.Dixie.Chat != "" || cfg.Search.AnyEnabled() || !cfg.Batch.Disabled {
|
||||
c, err := dcache.Open(cachePath, 500<<20, log)
|
||||
if err != nil {
|
||||
log.Warn("cache", "err", err)
|
||||
@@ -127,6 +129,23 @@ func serveMain() {
|
||||
cache = c
|
||||
}
|
||||
}
|
||||
if !cfg.Batch.Disabled && cache != nil {
|
||||
now := time.Now().Unix()
|
||||
if n := cache.MarkInFlightFailed("restart", now); n > 0 {
|
||||
for i := 0; i < n; i++ {
|
||||
srv.Prom.ObserveBatchJob("restart")
|
||||
}
|
||||
log.Info("batch reaped", "jobs", n)
|
||||
}
|
||||
srv.Batch = srv.NewBatchRunner(cache, batch.Config{
|
||||
TTL: config.ParseDuration(cfg.Batch.TTL, 0),
|
||||
TargetsPerJob: cfg.Batch.TargetsPerJob,
|
||||
JobsPerCaller: cfg.Batch.JobsPerCaller,
|
||||
JobsTotal: cfg.Batch.JobsTotal,
|
||||
Workers: cfg.Batch.Workers,
|
||||
MaxBytes: config.ParseSize(cfg.Batch.MaxBytes, 0),
|
||||
})
|
||||
}
|
||||
if cfg.Dixie.Chat != "" {
|
||||
srv.Digest = &digest.Engine{
|
||||
ChatURL: cfg.Dixie.Chat,
|
||||
@@ -213,8 +232,14 @@ func serveMain() {
|
||||
case <-sig:
|
||||
stop()
|
||||
_ = hs.Close()
|
||||
if srv.Batch != nil {
|
||||
srv.Batch.Close()
|
||||
}
|
||||
case err := <-errCh:
|
||||
stop()
|
||||
if srv.Batch != nil {
|
||||
srv.Batch.Close()
|
||||
}
|
||||
if err != nil {
|
||||
log.Error("serve", "err", err)
|
||||
os.Exit(1)
|
||||
|
||||
Vendored
+2
-1
@@ -28,7 +28,7 @@ RestrictRealtime=true
|
||||
RestrictSUIDSGID=true
|
||||
LockPersonality=true
|
||||
MemoryDenyWriteExecute=true
|
||||
MemoryMax=512M
|
||||
MemoryMax=1G
|
||||
SystemCallArchitectures=native
|
||||
SystemCallFilter=@system-service
|
||||
SystemCallFilter=~@privileged
|
||||
@@ -37,6 +37,7 @@ AmbientCapabilities=
|
||||
ReadWritePaths=/var/lib/sift
|
||||
ConfigurationDirectory=sift
|
||||
StateDirectory=sift
|
||||
EnvironmentFile=-/var/lib/sift/secret.env
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
|
||||
Vendored
+12
-7
@@ -1,5 +1,6 @@
|
||||
# sift configuration. P0 uses -listen and -deny-prefix; remaining keys are P1+.
|
||||
# This file is an example: the service does not assume Nix store paths.
|
||||
# sift configuration (/etc/sift/config.toml, or $SIFT_CONFIG). Every key is optional;
|
||||
# commented values are examples. The service does not assume Nix store paths.
|
||||
# Paid-search keys (BRAVE_API_KEY, TAVILY_API_KEY) go in /var/lib/sift/secret.env, not here.
|
||||
|
||||
listen = "127.0.0.1:8890"
|
||||
|
||||
@@ -8,11 +9,6 @@ listen = "127.0.0.1:8890"
|
||||
# Set on-box; do not commit residential addresses to git.
|
||||
# deny_prefixes = ["203.0.113.0/24"]
|
||||
|
||||
# [upstreams]
|
||||
# searxng = "http://127.0.0.1:8888"
|
||||
# dixie = "http://dixie.scylla-hammerhead.ts.net:11434/v1"
|
||||
# honcho = "http://rift.scylla-hammerhead.ts.net:8000"
|
||||
|
||||
# user_agent = "sift/dev"
|
||||
|
||||
# [self_ip]
|
||||
@@ -58,3 +54,12 @@ listen = "127.0.0.1:8890"
|
||||
# url = "https://api.tavily.com/search"
|
||||
# daily_credit_cap = 500
|
||||
# no_rerank = true
|
||||
|
||||
# [batch]
|
||||
# disabled = false
|
||||
# targets_per_job = 32
|
||||
# jobs_per_caller = 1
|
||||
# jobs_total = 8
|
||||
# workers = 1 # at most 2: each worker fetches one target at a time
|
||||
# ttl = "24h" # a job's results live this long after its last read or write
|
||||
# max_bytes = "64MiB" # stored results; the oldest finished jobs go first
|
||||
|
||||
+435
@@ -0,0 +1,435 @@
|
||||
# Sift P3.1 — search filters and multi-query merge
|
||||
|
||||
Drafted 2026-09-29 by talos from Kyle's ask to steal useful Tavily-cookbook
|
||||
patterns for sift (thread in Mattermost). Reviewer: Claude. Implementer: titan.
|
||||
|
||||
Amended 2026-09-29 after Claude's review of `7f4f0d5` (R1–R10) and of `73ff7b6`
|
||||
(E1–E3). Spec approved; titan implements this SHA.
|
||||
|
||||
P3 shipped `/v1/search` + `/v2/search` with SearXNG / Brave LLM Context / Tavily
|
||||
fan-in, canonical-URL dedupe, RRF, cache, and a daily Tavily cap. The Tavily
|
||||
client hardcodes `search_depth: "basic"` and sends no filters
|
||||
(`internal/search/tavily.go`). Native search accepts only `query`, `k`,
|
||||
`rerank`, `digest`.
|
||||
|
||||
This phase adds the cookbook's high-leverage search knobs without taking crawl,
|
||||
map, research, or `include_answer`.
|
||||
|
||||
## 1. Scope
|
||||
|
||||
In:
|
||||
|
||||
- Native filters on `POST /v1/search` and `sift search`.
|
||||
- Passing those filters through to Tavily when that source is enabled.
|
||||
- Honest SearXNG approximations where the engine can do it; **notes** (not
|
||||
source failure) where it cannot.
|
||||
- Post-merge domain enforcement in sift for restrict-mode include and exclude,
|
||||
for every source.
|
||||
- Multi-query search: several focused queries, one URL-deduped result list,
|
||||
fused with RRF across per-query lists.
|
||||
- Optional `min_score` cut on **Tavily's own score**, per source, before merge.
|
||||
- `published_date` on a hit when a source supplies it (Tavily news; SearXNG
|
||||
`publishedDate` when present).
|
||||
- Skill recipe: search → score/domain filter → `sift batch`, do not fetch every hit.
|
||||
|
||||
Out:
|
||||
|
||||
- `/v2/search` stays the Firecrawl contract (`query`, `limit` only). Unknown
|
||||
fields ignored. No filters on v2. v2 builds its own output map and must not
|
||||
leak `published_date`.
|
||||
- Crawl, map, Tavily `/extract` `/crawl` `/map`, `include_answer`,
|
||||
`include_raw_content`, `include_images` — P3 §3.3 still holds.
|
||||
- HTML lexical chunks, same-host map, Research/SSE/JSON-schema — later phases.
|
||||
- Changing `internal/guard/` or `internal/auth/`.
|
||||
- Fetching result URLs inside search.
|
||||
|
||||
Guard and contract: filters only change query text, SearXNG query params, and
|
||||
the Tavily JSON body. Endpoints stay config-built. Nothing from the request
|
||||
reaches a dialer. No guard or auth change.
|
||||
|
||||
## 2. Native request (`POST /v1/search`)
|
||||
|
||||
Existing fields unchanged. Additive:
|
||||
|
||||
```json
|
||||
{
|
||||
"query": "…",
|
||||
"queries": ["…", "…"],
|
||||
"k": 8,
|
||||
"rerank": true,
|
||||
"digest": false,
|
||||
"include_domains": ["openai.com"],
|
||||
"exclude_domains": ["pinterest.com"],
|
||||
"include_domains_mode": "restrict",
|
||||
"time_range": "month",
|
||||
"start_date": "2026-09-01",
|
||||
"end_date": "2026-09-29",
|
||||
"topic": "news",
|
||||
"exact_match": false,
|
||||
"min_score": 0
|
||||
}
|
||||
```
|
||||
|
||||
`time_range` and `start_date`/`end_date` are mutually exclusive (Tavily 400s if
|
||||
both are sent). The example above shows the fields, not a legal combination.
|
||||
|
||||
### 2.1 Error strings (tests compile against these)
|
||||
|
||||
- Query problems → 400, `"error": "invalid query"`, empty hint.
|
||||
- Filter problems → 400, `"error": "invalid search"`, `"hint"` names the field
|
||||
(`include_domains`, `exclude_domains`, `include_domains_mode`, `time_range`,
|
||||
`start_date`, `end_date`, `topic`, `min_score`, `exact_match`).
|
||||
|
||||
### 2.2 Validation
|
||||
|
||||
- Exactly one of `query` (non-empty) or `queries` (1–4 non-empty strings). Both
|
||||
set, or `queries` longer than 4, or any empty element → `invalid query`.
|
||||
Each string still has P3's 1500-rune cap (`invalid query`). ASCII `"` is
|
||||
allowed in queries (agents quote phrases themselves).
|
||||
- `exact_match: true` requires at least one balanced quoted phrase in **each**
|
||||
query (a `"`…`"` pair with one or more runes between; no unpaired `"`).
|
||||
Otherwise `invalid search` hint `exact_match`. Do not wrap or rewrite the
|
||||
query. `exact_match: false`/omitted: quotes are ordinary query text.
|
||||
- `include_domains` / `exclude_domains`: hostnames only. Max 20 include, 20
|
||||
exclude after normalise+dedupe. Empty list = omitted. Bad value →
|
||||
`invalid search` hint `include_domains` or `exclude_domains`.
|
||||
- Domain normalisation (apply before validation and before any `site:` rewrite):
|
||||
lowercase; strip a trailing dot; ASCII letters/digits/hyphens (punycode
|
||||
`xn--` labels allowed); no scheme, port, path, wildcard, or space; at most
|
||||
253 characters. Letters/digits/hyphens already make operator injection
|
||||
impossible (that needs a space). Do **not** substring-reject `OR` —
|
||||
`oracle.com` is valid. A space, e.g. `a.com OR b.com`, → 400.
|
||||
- `include_domains_mode`: omitted/`restrict`/`prefer`. Set without
|
||||
`include_domains` → `invalid search` hint `include_domains_mode`.
|
||||
Omitted with `include_domains` present means `restrict`.
|
||||
- `time_range`: omitted or `day`/`week`/`month`/`year` (and Tavily aliases
|
||||
`d`/`w`/`m`/`y`, normalised to the long form). Else `invalid search` hint
|
||||
`time_range`.
|
||||
- `start_date` / `end_date`: `YYYY-MM-DD` or omitted. `end` before `start` →
|
||||
`invalid search` hint `end_date`. Either date together with `time_range` →
|
||||
`invalid search` hint `time_range`.
|
||||
- `topic`: omitted/`general`/`news`/`finance`. Else `invalid search` hint `topic`.
|
||||
- `min_score`: omitted or `0` = off; otherwise `0 < x ≤ 1`. Else `invalid search`
|
||||
hint `min_score`.
|
||||
- `digest: true` still warns `digest ignored` (P3). That warning is API-level,
|
||||
not a source failure.
|
||||
|
||||
`/v2/search` does not grow. Hermes never sends these fields.
|
||||
|
||||
## 3. Response and source notes
|
||||
|
||||
`/v1/search` results gain optional `"published_date"` (RFC 1123 or ISO date
|
||||
string as the vendor sent it; omit the key when unknown). Do not invent dates.
|
||||
|
||||
`Hit` grows `PublishedDate string` (`json:"published_date,omitempty"`). When
|
||||
Dedupe or RRF replaces a hit's text with a longer snippet, keep the first
|
||||
non-empty `published_date`. Map SearXNG `publishedDate` when the engine
|
||||
supplies it.
|
||||
|
||||
`sizes` unchanged. `cached` unchanged.
|
||||
|
||||
### 3.1 `SourceOut.Notes` (R1 — do not reuse `Warn`)
|
||||
|
||||
Today a non-empty `SourceOut.Warn` means failure (`engine.go` Run loop): the
|
||||
source's hits are discarded, it lands in `sources_failed`, it is counted in
|
||||
`search_errors`, and caching is disabled. Filter honesty must not take that path.
|
||||
|
||||
```
|
||||
type SourceOut struct {
|
||||
Hits []Hit
|
||||
Warn string // failure only (timeout, HTTP error, pin)
|
||||
Notes []string // honesty: filters not applied, recall loss
|
||||
Credits int
|
||||
}
|
||||
```
|
||||
|
||||
- `Warn` keeps today's meaning. Hits discarded; source failed; no cache.
|
||||
- `Notes` are appended to the response `warnings` (deduplicated across queries).
|
||||
Hits are kept; the source is in `sources_ok`; the result may be cached.
|
||||
- Test: hits plus a note → hits kept, searxng in `sources_ok`, result cached.
|
||||
|
||||
## 4. Tavily client
|
||||
|
||||
Still `POST` the configured URL only. Still `include_answer: false`,
|
||||
`include_raw_content: false`, `include_images: false`, `search_depth: "basic"`.
|
||||
Still drop inbound `answer`. Still never call `/extract`, `/crawl`, `/map`.
|
||||
|
||||
Always send `auto_parameters: false`. Auto mode can choose `advanced`, which
|
||||
costs 2 credits and breaks "credit cost is fixed".
|
||||
|
||||
When the request has filters, add only the fields that are set:
|
||||
|
||||
- `include_domains`, `exclude_domains`
|
||||
- `include_domains_mode` **only when** `include_domains` is set
|
||||
- `time_range` **or** `start_date`/`end_date`, never both
|
||||
- `topic` when not `general` (omitted `topic` is `general`)
|
||||
- `exact_match: true` when the flag is set (query text unchanged, including
|
||||
the agent's quotes)
|
||||
- `include_published_date: true` when `topic` is `news` or a date filter is set
|
||||
|
||||
Leave `filter_by_published_date` unset. Tavily's `time_range` and dates
|
||||
filter on publish date **or last updated date**; `filter_by_published_date`
|
||||
also drops undated results. Unset means undated and recently updated pages
|
||||
can pass the date filter.
|
||||
|
||||
Leave `include_usage` unset. `usage.credits` only comes back when requested;
|
||||
today's count is the code's fallback of 1, which is correct for `basic`.
|
||||
|
||||
Map `results[].published_date` onto `Hit`. Count `usage.credits` as today
|
||||
(fallback 1).
|
||||
|
||||
A fixture test must still prove `include_answer` is false, `auto_parameters` is
|
||||
false, `search_depth` is `basic`, and no extract/crawl path exists.
|
||||
|
||||
## 5. SearXNG and Brave
|
||||
|
||||
Do not pretend they support Tavily's full filter set. Vendor gaps are **notes**,
|
||||
never `Warn`. Domain restrict/exclude are enforced in sift after merge (§6.3);
|
||||
`site:` is a recall optimisation only.
|
||||
|
||||
**SearXNG** (`GET {url}/search`, same client/dialer as P3):
|
||||
|
||||
- One `include_domains` entry: prefix the query with `site:<host> `.
|
||||
- Several: `site:a OR site:b OR …` up to 3, then note
|
||||
`searxng: extra include_domains dropped`.
|
||||
- `exclude_domains`: no query rewrite (post-filter enforces). No note unless
|
||||
recall is otherwise lost.
|
||||
- `include_domains_mode=prefer`: no `site:` prefix, no post-filter; note
|
||||
`searxng: include_domains_mode=prefer not applied`.
|
||||
- `time_range`: send SearXNG query param `time_range=` (`day`/`week`/`month`/`year`).
|
||||
Never a bang in the query text.
|
||||
- `topic=news`: send `categories=news`. Never `!news`.
|
||||
- `topic=finance`: note `searxng: topic ignored` (no finance category).
|
||||
- `exact_match`: send the query unchanged. SearXNG engines already treat
|
||||
quotes as phrases. Do not wrap again.
|
||||
- `start_date`/`end_date`: ignored + note `searxng: start_date/end_date not applied`.
|
||||
- `min_score > 0`: note `searxng: min_score not applied`; keep hits.
|
||||
|
||||
Live instance (straylight, 2026-09-29, not independently verified here):
|
||||
`time_range=month` and `categories=news` both returned 200 with normal result
|
||||
lists. Send both params. Tests are offline fixtures.
|
||||
|
||||
**Brave LLM Context**: no documented domain/date/topic/exact filters. Do not
|
||||
add a second Brave HTTP API. Post-filter still enforces restrict include and
|
||||
exclude. Notes when the corresponding filter was set:
|
||||
|
||||
- `brave_llm: time_range not applied`
|
||||
- `brave_llm: topic not applied`
|
||||
- `brave_llm: exact_match not applied`
|
||||
- `brave_llm: include_domains_mode=prefer not applied`
|
||||
- `brave_llm: min_score not applied`
|
||||
|
||||
Request URL unchanged aside from `q` (do not wrap for `exact_match`; Brave
|
||||
cannot apply it). Deduplicate identical notes across queries.
|
||||
|
||||
## 6. Pipeline
|
||||
|
||||
`Source.Search` grows a `Filter` argument (zero value = today's behaviour).
|
||||
|
||||
### 6.1 Extra queries run sequentially per source (R5)
|
||||
|
||||
Sources still run in parallel **within** one query, as today. Queries run
|
||||
**sequentially**. Up to 4 queries means up to 4 Brave and 4 Tavily paid calls;
|
||||
four parallel Brave calls will hit the plan's rate limit, and a 429 marks
|
||||
brave failed (`Warn`) and disables caching.
|
||||
|
||||
A Tavily call still counts toward the daily cap. If the cap trips mid-batch,
|
||||
later queries skip Tavily with the same `tavily: daily cap` warning (skip path,
|
||||
not a failed attempted source). The queries then ran against different source
|
||||
sets. **Do not cache** that result (the key's `sources` would be wrong). Test
|
||||
this.
|
||||
|
||||
### 6.2 `min_score` (R2)
|
||||
|
||||
After merge, `Hit.Score` is the RRF value `Σ 1/(60+rank)`. With 3 sources the
|
||||
maximum is about 0.05, so a post-merge `--min-score 0.5` would drop everything.
|
||||
With `rerank:false` scores are vendor-native (Tavily 0–1, SearXNG unbounded,
|
||||
Brave 0). With Tavily `no_rerank`, Tavily scores mix with RRF.
|
||||
|
||||
Apply `min_score` **per source, before the merge, to Tavily's own relevance
|
||||
score only**. That is the cookbook's usage. Other sources get a note
|
||||
(`searxng: min_score not applied`, `brave_llm: min_score not applied`) and keep
|
||||
their hits. Do not drop score-less Brave/SearXNG hits.
|
||||
|
||||
### 6.3 Domain post-filter (R3)
|
||||
|
||||
After each query's merge, for **every** source, drop hits that violate:
|
||||
|
||||
- `include_domains` when mode is `restrict` (omitted mode = restrict)
|
||||
- `exclude_domains` (always)
|
||||
|
||||
Match is exact on the result host: `host == d` or `host` has suffix `"." + d`
|
||||
(host lowercased, trailing dot stripped). Cost is negligible.
|
||||
|
||||
- `prefer` mode: **no** post-filter. Tavily honours `include_domains_mode`;
|
||||
other sources get the prefer note.
|
||||
- `site:` is a recall optimisation only. Notes are for recall loss, not for
|
||||
"filter ignored" when the post-filter already enforced it.
|
||||
- As written without this, Brave ignores the filter entirely and SearXNG
|
||||
engines that do not understand `site:` return off-domain hits.
|
||||
|
||||
### 6.4 Multi-query merge (R4)
|
||||
|
||||
Do not concatenate-then-truncate: if q1 returns `k` or more hits, q2–q4
|
||||
contribute nothing.
|
||||
|
||||
For `queries: [q1, q2, …]` (max 4), or a single `query` treated as a list of one:
|
||||
|
||||
1. For each query, sequentially: existing fan-in (parallel sources, same
|
||||
filters, same per-source timeout) → Tavily `min_score` → today's
|
||||
`mergeSearch` (RRF / no-rerank path) → domain post-filter.
|
||||
2. **RRF across the per-query fused lists** (each query's list is one ranked
|
||||
list). URLs several queries found rise.
|
||||
3. Dedupe (P3: first source wins attribution; later longer snippet replaces
|
||||
text; keep first non-empty `published_date`).
|
||||
4. Truncate to `k`.
|
||||
|
||||
Test: q1 returns `k` or more unique hits, and q2's single distinct top hit
|
||||
still appears.
|
||||
|
||||
Do not fetch any result URL.
|
||||
|
||||
## 7. Cache
|
||||
|
||||
P3 key is `sha256(norm(query) | sources | k | hour_bucket | caller_class)`.
|
||||
|
||||
Replace the payload with a canonical JSON object, keys sorted, then sha256.
|
||||
It must hold:
|
||||
|
||||
- `caller_class`, hour bucket (`2006-01-02T15` UTC), `k`, attempted `sources`
|
||||
(sorted) — exactly as P3 has them
|
||||
- `queries`: normalised, ordered list (a single `query` is a one-element list)
|
||||
- `include_domains` / `exclude_domains`: normalised, sorted, deduplicated
|
||||
- `include_domains_mode` (`restrict` or `prefer`; omit the key when no include
|
||||
list)
|
||||
- `time_range` in long form (`d` → `day`); dates as `YYYY-MM-DD` or omit
|
||||
- `topic`, with omitted meaning `general` so the two share a key
|
||||
- `exact_match` (bool)
|
||||
- `min_score` in one fixed number format (e.g. `'%g'` is not stable — use
|
||||
`strconv.FormatFloat(x, 'f', -1, 64)` or a documented equivalent, same for
|
||||
0 and omitted)
|
||||
|
||||
Same query without filters must not hit a filtered entry.
|
||||
|
||||
Still: do not cache if any attempted source failed (`Warn`). Still do not cache
|
||||
if the Tavily cap tripped partway through a multi-query run. Still no page
|
||||
bodies, no Tavily answers.
|
||||
|
||||
## 8. Logs
|
||||
|
||||
P3 §8 already bans hosts in the search log for every caller, and filter values
|
||||
are hosts or query content.
|
||||
|
||||
- Never log domain values, query text, dates, or topic values.
|
||||
- Restricted callers keep **exactly** P3's fields: `caller`, `q_len`, `sources`,
|
||||
`n`, `cache`.
|
||||
- Fleet callers may add `n_queries` and the filter *kinds*
|
||||
(`filters=domains,time,topic`), never values. `rerank_ms` stays as today.
|
||||
- `q_len` for several queries: sum of their rune lengths.
|
||||
|
||||
Test: restricted and fleet requests with `include_domains: ["secret.example"]`
|
||||
and 3 queries → no log record contains the host or any query.
|
||||
|
||||
## 9. CLI
|
||||
|
||||
```
|
||||
sift search "<query>" [--k n] [--json]
|
||||
[--include-domain host]... [--exclude-domain host]...
|
||||
[--domains-mode restrict|prefer]
|
||||
[--time-range day|week|month|year]
|
||||
[--start-date YYYY-MM-DD] [--end-date YYYY-MM-DD]
|
||||
[--topic general|news|finance]
|
||||
[--exact-match]
|
||||
[--min-score 0.5]
|
||||
[--queries-file path]
|
||||
```
|
||||
|
||||
`--min-score 0.5` is Tavily's 0–1 scale (applied only to Tavily, §6.2).
|
||||
|
||||
`--queries-file` is a JSON array of 1–4 strings. It is mutually exclusive with
|
||||
the positional query. Repeatable `--include-domain` / `--exclude-domain`.
|
||||
|
||||
Default text output: if `published_date` is present, print it as a fourth line
|
||||
under the snippet. `--json` is the native envelope.
|
||||
|
||||
## 10. Skill
|
||||
|
||||
`skill/SKILL.md`:
|
||||
|
||||
- Document the new flags with one example matching the cookbook Product News
|
||||
Tracker: official domain search + `topic=news` + `time_range=month`.
|
||||
- **Recipe:** after `sift search --json`, do not `sift fetch` every URL. Drop
|
||||
hits below a Tavily score you care about (or take top N), then `sift batch submit`.
|
||||
- `web_search` remains the fallback only when sift is unreachable. Hermes
|
||||
`web_search` still hits `/v2/search` and does not get filters — say so.
|
||||
|
||||
## 11. Tests (minimum)
|
||||
|
||||
Owner-given invariants (these are where a free-form implementation drifts):
|
||||
|
||||
- **R1 notes:** hits plus a note → hits kept, source in `sources_ok`, result
|
||||
cached. A `Warn` still fails the source and disables cache.
|
||||
- **R2 min_score:** Tavily hit at 0.4 dropped when `min_score=0.5`; SearXNG /
|
||||
Brave hits kept; notes present. Post-merge RRF scores are not compared to
|
||||
`min_score`.
|
||||
- **R3 domain post-filter:** Brave (and SearXNG) hit whose host is not in
|
||||
restrict-include, or is in exclude, is dropped. `prefer` keeps it.
|
||||
Invalid domain with a space, e.g. `a.com OR b.com`, → 400 `invalid search`.
|
||||
`oracle.com` is valid.
|
||||
- **R4 multi-query RRF:** q1 returns k unique hits; q2's distinct top hit still
|
||||
appears.
|
||||
- **R5 cap mid-batch:** Tavily cap trips after query 1 → later queries skip
|
||||
Tavily; result is not cached.
|
||||
- **R6 cache key:** same query different `include_domains` → miss; same
|
||||
filters → hit; omitted `topic` and `topic=general` share a key; `d` and
|
||||
`day` share a key.
|
||||
- **R9 logs:** restricted and fleet with `include_domains:["secret.example"]`
|
||||
and 3 queries → no log record contains the host or any query.
|
||||
- Tavily outbound JSON: filters present when set; `include_answer` always
|
||||
false; `auto_parameters` false; `search_depth` basic; no extract/crawl;
|
||||
`include_domains_mode` absent unless include list set; never both
|
||||
`time_range` and dates; `exact_match: true` with query containing `"phrase"`
|
||||
→ flag true and query unchanged (still quoted); `exact_match: true` with no
|
||||
quoted phrase → 400 `invalid search` hint `exact_match`.
|
||||
- Tavily inbound: `answer` dropped; `published_date` mapped; missing date omitted.
|
||||
- Validation table for §2 (`invalid query` vs `invalid search` + hint).
|
||||
- SearXNG: `site:` prefix for one domain; `exact_match` does not wrap `q`
|
||||
(query already quoted by the agent); `time_range` / `categories` as query
|
||||
params not bangs.
|
||||
- Brave: request URL unchanged (aside from `q`); notes as §5; post-filter still
|
||||
drops off-domain hits.
|
||||
- `/v2/search` golden unchanged (no new keys required, no `published_date`).
|
||||
- Restricted caller: same sources; logs still have no query/host/URL.
|
||||
|
||||
`make gate` (and contract if P3's contract harness is still in the tree).
|
||||
|
||||
## 12. Code layout
|
||||
|
||||
Touch, do not explode:
|
||||
|
||||
- `internal/search` — `Filter` type; `SourceOut.Notes`; extend `Source.Search`;
|
||||
Tavily/SearXNG/Brave; engine loop; domain post-filter; multi-query RRF
|
||||
- `internal/api/search.go` — native decode + validation + log fields
|
||||
- `cmd/sift/searchcmd.go` — flags
|
||||
- `internal/dcache` search key
|
||||
- `skill/SKILL.md`, `README.md`
|
||||
|
||||
No new dependency. No `internal/api/api.go` growth if a new file is cleaner
|
||||
(`internal/api/search_filter.go`). Guard and auth stay closed.
|
||||
|
||||
## 13. Rollout
|
||||
|
||||
1. Claude approved `73ff7b6` once E1–E3 are in (this SHA). No further spec round.
|
||||
2. SearXNG probe already returned 200 for `time_range=` and `categories=news`.
|
||||
3. titan implements on `git.wntrmute.dev:kyle/sift` from this SHA. Split commits
|
||||
if it helps; one phase doc (invariants in §11 — titan writes those tests;
|
||||
Claude checks them in the code review). titan does not edit
|
||||
`internal/guard/`, `internal/auth/`, or `~/src/nixos`.
|
||||
4. Claude reviews the code.
|
||||
5. straylight bumps the orion pin.
|
||||
|
||||
## 14. Explicit non-goals (do not sneak in)
|
||||
|
||||
Map, crawl, HTML BM25 chunks, Research, `include_answer`, social-platform
|
||||
crawlers, changing `/v2/search`. Those were ranked later in the cookbook
|
||||
write-up on purpose.
|
||||
+268
@@ -0,0 +1,268 @@
|
||||
# Sift P6 — asynchronous batch fetch
|
||||
|
||||
Drafted 2026-09-26 by straylight (Hermes agent) from the design discussion with Kyle. Reviewer:
|
||||
Claude (when it frees up this evening). Implementer: Kyle's call.
|
||||
|
||||
## 1. Scope
|
||||
|
||||
In: an asynchronous batch fetch endpoint. An agent submits a list of targets (URLs, each with the
|
||||
same per-target options as `/v1/fetch`) and gets a job ID back immediately; it polls for per-target
|
||||
results later. One job drains **serially** — one target at a time — in the background.
|
||||
|
||||
Out:
|
||||
|
||||
- Parallel fan-out. A batch is deliberately slow; it is the tool for research threads that are
|
||||
already expected to take a long time, not for interactive fetches.
|
||||
- Crawling, retry-with-backoff at the target level, and anything that requires a headless browser.
|
||||
- Firecrawl compatibility. This is a native-only feature; no Hermes plugin.
|
||||
|
||||
The guard, auth and restricted-caller rules do not change: each target is fetched exactly like a
|
||||
page, through the same SSRF guard and the same per-caller/host gates.
|
||||
|
||||
## 2. Why agents need this (2026-09-26)
|
||||
|
||||
`docs/PLAN.md` §13 (Kyle): "an agent that has collected a series of targets and wants them fetched
|
||||
without holding a request open per target." Today that means a loop of `sift fetch` calls — one
|
||||
request held open per URL, each its own tool call and gate. A batch lets the agent submit the whole
|
||||
list, walk away, and collect once. It builds on the P5 infrastructure that already exists: the
|
||||
`dcache` SQLite store, the `global`/`per` concurrency gates, and the `doFetch` pipeline.
|
||||
|
||||
## 3. Endpoints and wire shapes
|
||||
|
||||
| Method | Path | Result |
|
||||
|---|---|---|
|
||||
| `POST /v1/batch` | submit a list of targets | `202` created/reused, `400` malformed, `422` refused up front or over a cap, `429` caller/jobs busy |
|
||||
| `GET /v1/batch/{id}` | status summary; `?results=1` inlines results | `200`, `401`/`404` |
|
||||
| `GET /v1/batch/{id}/{idx}` | one target's status and result | `200`, `401`/`404` |
|
||||
| `DELETE /v1/batch/{id}` | cancel | `200`, `409` already finished, `401`/`404` |
|
||||
|
||||
With `[batch] disabled = true` every batch endpoint answers `503 batch disabled`.
|
||||
|
||||
`POST /v1/batch` request — each target is the `/v1/fetch` field set minus the outer URL:
|
||||
|
||||
```json
|
||||
{"targets": [{"url": "…", "mode": "clean", "question": "…", "pages": "…", "grep": "…",
|
||||
"jq": "…", "context": true, "recover": false}, …]}
|
||||
```
|
||||
|
||||
Response `202`:
|
||||
```json
|
||||
{"success": true, "job_id": "…", "status": "queued", "targets": 3, "reused": false}
|
||||
```
|
||||
|
||||
`GET /v1/batch/{id}` (a summary: no results, since 32 inline results can be ~6 MB in one agent
|
||||
tool call):
|
||||
```json
|
||||
{"job_id": "…", "status": "partial", "targets_total": 3, "targets_done": 3,
|
||||
"created_at": "2026-09-27T01:02:03Z", "targets": [
|
||||
{"idx": 0, "url": "…", "status": "done"},
|
||||
{"idx": 1, "url": "…", "status": "failed", "error": "blocked: private address …", "hint": "browser"}
|
||||
]}
|
||||
```
|
||||
|
||||
`?results=1` adds `"result"` (the `/v1/fetch` envelope) to every `done` target.
|
||||
`GET /v1/batch/{id}/{idx}` returns `{"job_id", "idx", "url", "status", "result"?, "error"?,
|
||||
"hint"?}`. A `pending` target has `status: "pending"` and no `result`/`error`. A `202` for a new
|
||||
job says `"status": "queued"`; a reused one says the job's current status. `DELETE` → `200 {"success": true, "status":
|
||||
"cancelled"}`.
|
||||
|
||||
Job status vocabulary: `queued → running → done | partial | failed | cancelled`. Target status:
|
||||
`pending → running → done | failed`.
|
||||
|
||||
## 4. Storage (`dcache`)
|
||||
|
||||
Two tables, added to `internal/dcache` (which owns the single SQLite connection):
|
||||
|
||||
- `batch_job(id, caller, caller_class, idem_key, status, error, targets_total, targets_done,
|
||||
created_at, lease, nbytes)`
|
||||
- `batch_target(job_id, idx, url, options, status, result, error, hint)` — `result` is the JSON
|
||||
envelope, `NULL` until the target finishes.
|
||||
|
||||
Batch storage is a coherent whole: it is bounded by its own lease GC and a byte budget, **not** by
|
||||
the content cache's LRU. `EvictBatchOldest` evicts whole jobs oldest-`lease`-first; a
|
||||
`queued`/`running` job is never byte-evicted. `MarkInFlightFailed(reason, now)` is the restart hook.
|
||||
|
||||
## 5. The runner (`internal/batch`)
|
||||
|
||||
A new package with the job lifecycle and the serial state machine. Persistence is the `*dcache.DB`;
|
||||
the only injected behavior is a `FetchFunc func(ctx, Target) (Result, error)`.
|
||||
|
||||
- **Idempotency (in v1).** A job has a random opaque `id` plus a derived `idem_key =
|
||||
SHA256(caller + "\x00" + canonical target list)`. Resubmitting the same (caller, targets) while
|
||||
the job's lease has not expired returns the existing job (`reused: true`) instead of a new one.
|
||||
Idempotency is lease-scoped, so it is "reuse within TTL, fresh after".
|
||||
- **Caps.** targets/job 32; live jobs/caller 1; live jobs total 8. "Live" is by status
|
||||
(`queued`/`running`), not lease. Exceeding a cap is a typed error the API maps to `422`/`429`.
|
||||
- **The caller.** `Submit` takes the `*auth.Caller`; the runner keeps it in memory with the job
|
||||
and passes it to every `FetchFunc` call (`func(ctx, *auth.Caller, Target)`), so the fetch runs
|
||||
under the caller's gates, pool slot, cache class and redaction. Only its name and class are
|
||||
stored.
|
||||
- **Workers.** `workers` (default 1, at most 2) share one round-robin over the live jobs, one
|
||||
target at a time. A job is `queued` until a worker first takes it, then `running`.
|
||||
- **Serial.** Targets drain in `idx` order; target `i+1` does not start until `i` is written.
|
||||
- **Busy.** A `FetchFunc` error wrapping `ErrRetryLater` (the API maps fetch `429`/`503` to it)
|
||||
leaves the target pending; it is tried again after 2 s, for up to 10 min, then failed.
|
||||
- **Failure isolation.** One target's error is stored per-target and the job continues; the job
|
||||
status is `done` (all ok), `failed` (all failed), or `partial` (mixed).
|
||||
- **Cancel.** `DELETE` flips the job `cancelled`, stops before the next target, and lets the
|
||||
in-flight target finish and store. Completed targets stay readable until the lease lapses.
|
||||
|
||||
## 6. Scheduling and fairness
|
||||
|
||||
The workers run in the background, and every target still goes through the normal `acquire` gate
|
||||
(8 global, 2 per-host). With 1–2 workers the batch holds at most 1–2 of the 8 global slots, so
|
||||
interactive requests always have room. (The first draft let 8 live jobs each run a target at
|
||||
once, which could fill the whole gate.) On `429 busy` the worker **waits and retries** instead of
|
||||
failing: nobody is holding a socket open.
|
||||
|
||||
Two shared resources get the same treatment:
|
||||
- **The PDF pool.** A batch extraction uses the pool caller name `<caller>/batch`, so it never
|
||||
takes the caller's own interactive slot. The single-flight key is unchanged, so a batch and an
|
||||
interactive request for the same PDF still share one child.
|
||||
- **The digest gate.** A batch question queues on it like any other digest. With one worker
|
||||
that is at most one batch digest at a time.
|
||||
|
||||
Each batch fetch runs under the same 210 s deadline as a `/v1/fetch` request.
|
||||
|
||||
## 7. Retention and TTL
|
||||
|
||||
A job carries a single sliding **lease**, bumped on submit, on every target completion, and on
|
||||
every read by its owner (the API calls `Touch` after the ownership check; the store's reads
|
||||
never write). All targets therefore share one expiry (the last completion's), and a
|
||||
slow consumer's reads keep a finished batch alive. GC deletes a job and its targets once the lease
|
||||
is older than the TTL (24 h). Completed results survive a restart and stay readable until the lease
|
||||
lapses; `queued`/`running` jobs are marked `failed` (`restart`) on startup and never resumed.
|
||||
|
||||
## 8. Restricted callers
|
||||
|
||||
Compartmentalized-access systems: their usage is redacted as much as possible. The existing
|
||||
`logHost`/`logReason`/`logErr` machinery already blanks host, reason and error detail for
|
||||
`PermitRestricted` callers, and the digest/embed caches partition by caller class — the batch
|
||||
inherits all of it by passing the caller through. Two additions: batch log lines never carry a
|
||||
target URL or hostname for **any** caller, and `GET`/`DELETE` enforce job ownership, returning
|
||||
`404` (not `403`) for a job owned by another caller so a caller cannot probe for its existence.
|
||||
The `idem_key` is stored as a hash and never logged alongside its target list.
|
||||
|
||||
## 9. Caps and metrics (sized empirically)
|
||||
|
||||
The caps are starting values, instrumented so they can be re-sized from observed data:
|
||||
|
||||
- `sift_batch_jobs_total{outcome}` — `created|reused|done|partial|failed|cancelled|restart`
|
||||
- `sift_batch_targets_total{outcome}` — `ok|blocked|timeout|busy|upstream_4xx|upstream_5xx|
|
||||
unsupported|too_long|cancel`
|
||||
- `sift_batch_queue_depth`, `sift_batch_targets_inflight` (gauges)
|
||||
- `sift_batch_job_seconds`, `sift_batch_target_seconds`, `sift_batch_result_bytes` (histograms)
|
||||
|
||||
Both `outcome` counters fold any value outside their known set to `other`, so a raw string never
|
||||
becomes a label. `result_bytes` is what sizes the byte budget (64 MiB).
|
||||
|
||||
## 10. CLI and skill
|
||||
|
||||
- `sift batch submit <targets.json>` (`-` reads stdin; prints `job_id` and `reused`).
|
||||
- `sift batch status [--results] <job_id>` (the summary; `--results` inlines every result).
|
||||
- `sift batch get <job_id> <idx>` (one result).
|
||||
- `sift batch cancel <job_id>`.
|
||||
- `skill/SKILL.md`: use `sift batch submit` for a list of targets instead of a loop of
|
||||
`sift fetch`; `sift batch status` to collect.
|
||||
- `[batch]` config block: `disabled`, `targets_per_job`, `jobs_per_caller`, `jobs_total`,
|
||||
`workers`, `ttl`, `max_bytes`.
|
||||
The `/v1/batch` body cap (256 KiB, vs the 64 KiB global) is a constant, not config.
|
||||
|
||||
## 11. Refactor of existing code (explicit)
|
||||
|
||||
`v1fetch` today inlines the whole per-target pipeline (guard → `doFetch` → `recover` →
|
||||
render/digest → context → envelope). The batch worker needs the same pipeline for one target. It is
|
||||
extracted into one shared method, `func (s *Server) fetchTarget(ctx, *auth.Caller, batch.Target)
|
||||
(batch.Result, error)`, that both `v1fetch` and the runner's adapter (`batchFetch`) call. It is its
|
||||
own task (04), done before any batch endpoint exists, so the diff is a pure move. `v1fetch` keeps its
|
||||
exact response shape and behavior — its existing tests must pass unchanged. This is the one change
|
||||
to working code in P6 and is called out for reviewer attention.
|
||||
|
||||
## 12. Acceptance (contract) tests
|
||||
|
||||
Acceptance-first, as in P5a: given tests compiled against a panic-only skeleton, failing for the
|
||||
intended reasons. See `docs/plans/p6/_files/` and the task table in `docs/plans/p6/README.md`.
|
||||
|
||||
1. `dcache`: round trip, idempotency lookup (lease expiry, caller scope, status filter), live
|
||||
counts, ordered targets, exactly-once `targets_done` and lease bumps, sweep, byte eviction
|
||||
(oldest first, never a live job), `MarkInFlightFailed`, `BatchNbytes`.
|
||||
2. `batch` runner:
|
||||
- serial order (peak in flight 1), with the caller passed to every fetch;
|
||||
- round-robin between jobs, and `queued` → `running`;
|
||||
- failure isolation (`partial`, `failed`);
|
||||
- idempotent resubmit, including after `done`, but a new job after a cancel;
|
||||
- all three caps and `ErrNoTargets`; cancel (the in-flight target stored, later ones `pending`);
|
||||
- busy retry and its limit; panic recovery (the panic value not echoed); the hooks.
|
||||
3. `metrics`: series names, labels, folding to `other`.
|
||||
4. `api` `fetchTarget`: the same content and envelope as `/v1/fetch`; its errors as `*fetch.Error`.
|
||||
5. `api` wiring: the runner through `fetchTarget`, a private-resolving name blocked per target,
|
||||
metrics through the hooks, outcome classes, busy → `ErrRetryLater`.
|
||||
6. `guard.PreCheck`: no resolution; the port, scheme, userinfo and literal rules; the allow list.
|
||||
7. `api` endpoints:
|
||||
- submit `202` and idempotent resubmit;
|
||||
- `400`/`422` refusals (the guard pre-check naming the index, 33 targets);
|
||||
- the 256 KiB body; the per-target `blocked` failure giving `partial`;
|
||||
- summary, one result, `?results=1`;
|
||||
- ownership `404` (GET, one result, DELETE); cancel and `409`;
|
||||
- no job id in `/metrics`.
|
||||
8. `config`: the `[batch]` block, `ParseSize`. `cmd/sift`: `parseBatchArgs`.
|
||||
|
||||
## 13. Rollout
|
||||
|
||||
1. Implementer builds P6 on a branch; the reviewer (Claude) checks §11, the `v1fetch` refactor,
|
||||
and the caps/lease semantics.
|
||||
2. Orion pin bump with the `[batch]` config block; no `MemoryMax` change.
|
||||
3. Skill update ships with the pin bump.
|
||||
4. Watch `sift_batch_*` and tirith for a week; the success measure is that a research thread's
|
||||
one-request-per-URL fetch loop collapses into a single batch submit.
|
||||
|
||||
## 14. Changes after the plan was written
|
||||
|
||||
**2026-09-26 evening, owner review (Claude), at Kyle's direction ("fold them into the plan").**
|
||||
The design stands; these were implementation-blocking or would have hurt agents:
|
||||
|
||||
1. **The caller was lost.** `FetchFunc` had no caller and the runner was built once, so the given
|
||||
api test (`batch.New(db, s.fetchTarget, …)`) could not compile against the specified
|
||||
`fetchTarget(ctx, *auth.Caller, Target)`. Made to compile, it would have fetched every target
|
||||
with no caller: no per-caller gates, pool limit, cache class or redaction. Now `Submit` takes
|
||||
the `*auth.Caller` and `FetchFunc` receives it (§5).
|
||||
2. **Task 01 rule 1 contradicted its test.** Store reads no longer bump the lease; the API calls
|
||||
`Touch` after an owner's read (§7).
|
||||
3. **Batch could starve interactive fetches.** 8 live jobs × 1 target filled the 8-slot global
|
||||
gate. The busy-retry of §6 was in no task. Now there are 1–2 round-robin workers, busy-retry
|
||||
is in the runner, a batch PDF uses its own pool name, and a per-target deadline applies (§5, §6).
|
||||
4. **Flaky, networked API tests.** Two tests fetched `https://example.com/`, and the cancel test
|
||||
raced the fetch. Every api test is now local, using held `httptest` servers and a fake DNS
|
||||
name for the private-address case.
|
||||
5. **A resubmit could return a dead job.** A failed or cancelled job was reused for 24 h, and
|
||||
polling it kept extending that. Now only `queued`/`running`/`done`/`partial` jobs are reused.
|
||||
6. **Reads were too large.** `GET` inlined every result (up to ~6 MB). It is now a summary, with
|
||||
`GET …/{idx}` and `?results=1` (§3). The CLI gains `get` and `status --results`.
|
||||
7. **Body cap.** The global 64 KiB `MaxBytesReader` runs in `ServeHTTP`, so a handler cannot
|
||||
raise it. `ServeHTTP` now picks 256 KiB for `/v1/batch`. Request metrics use a fixed
|
||||
`/v1/batch/{id}` label, so a job id never becomes a label value.
|
||||
8. **Submit pre-check.** `guard.ParseURL` rejects the test's own loopback port and knows nothing
|
||||
of the allow list. New `Guard.PreCheck` is `CheckURL` without resolution (an owner commit on
|
||||
master, since `AGENTS.md` keeps the implementer out of `internal/guard/`). A private IP literal is now refused at submit; a *name* that resolves privately
|
||||
still fails per target at fetch time.
|
||||
9. **Task shape.** 5 tasks → 9 (the last split came during the run, below), one package each where possible:
|
||||
01 dcache, 02 runner core, 03 runner workers, 04 metrics, 05 the `fetchTarget` refactor alone,
|
||||
06 runner wiring, 07 endpoints, 08 serve/config, 09 CLI/docs. Every task restates the APIs it
|
||||
uses from earlier tasks. Every given test compiles against a panic-only skeleton and fails
|
||||
there on `panic: todo`.
|
||||
10. **Also fixed:**
|
||||
- the `ErrBusy` (jobs total) cap now has a test;
|
||||
- the `202` status is `queued`, consistent with the workers;
|
||||
- the `restart` outcome is counted by `serve`;
|
||||
- `Options.Context` is `*bool`, keeping `/v1/fetch`'s default.
|
||||
|
||||
**Sequencing.** P6 starts after P5c merges. P5b/P5c change `v1fetch` (the digest meta gains
|
||||
`retrieval`/`embedded`), which is the code task 05 moves. Recheck the api given tests against
|
||||
a skeleton built on post-P5c master before the run.
|
||||
|
||||
**During the run (2026-09-26 21:08).** Task 02 (the whole runner) ended with no file: the model
|
||||
wrote the implementation out in its reply, the session was compacted, and it then emitted a
|
||||
malformed tool call. The runner is now two tasks (02 core, 03 workers), each with its own given
|
||||
test. `AGENTS.md` gains a rule to write code into files early. Copy paths are spelled from the
|
||||
repo root. Details in `docs/plans/p6/README.md`, "Changes during the run".
|
||||
|
||||
@@ -177,6 +177,9 @@ approval prompt. Hence:
|
||||
setting (`user_agent`) overrides it for the whole service; no per-request UA and no
|
||||
automatic browser-UA retry. Revisit only if bot walls turn out to be common — measure it:
|
||||
count `upstream 403/429 (likely bot-blocked)` results per host in `/metrics`.
|
||||
**Revisit (Kyle, 2026-10-02):** UA spoofing is not off the table. Measurement that day
|
||||
(Reuters/Libération, DataDome) is below — a `User-Agent` change alone did not fetch
|
||||
those pages.
|
||||
|
||||
Phase placement: CLI `fetch --mode clean` + skill in **P1**; digest/context flags in P2;
|
||||
`search` in P3; `--recover` in P2 (it is what removes the titan-style prompts).
|
||||
@@ -310,3 +313,89 @@ harness); commit + push to `git.wntrmute.dev:kyle/sift`; NixOS changes go
|
||||
through `~/src/nixos` committed master and `rebuild-nixos` (fetch before editing — clones go
|
||||
stale); state directories are never deleted; secrets never printed; one phase at a time with
|
||||
its acceptance evidence posted in `~agent`.
|
||||
|
||||
## 13. Future work
|
||||
|
||||
### Digest retrieval misses terse summary facts (found 2026-09-26, after P5b)
|
||||
|
||||
On the RP2040 datasheet, "how many PIO state machines in total?" came back **4** from both
|
||||
`bm25` and `hybrid` retrieval. The answer, **8**, is a one-line bullet on p. 9 ("8 PIO state
|
||||
machines") that neither ranker put in the top 5 chunks, so the model answered from the
|
||||
per-block sections, and in the hybrid case also made up a split. `--grep` finds it at once, and
|
||||
the skill now says to confirm numbers with `--grep` (`94da9d9`).
|
||||
|
||||
Candidates, when there is a phase for it:
|
||||
- pin the chunk holding the document's feature or overview list, as chunk 0 is pinned today;
|
||||
- run a `--grep`-style exact-term pass over the question's numbers and nouns before ranking;
|
||||
- a small eval set of fact questions over the pdfbench corpus, to measure a change rather than guess.
|
||||
|
||||
### Paywalled pages → archive snapshot (Kyle, 2026-09-26; possible P7)
|
||||
|
||||
"If a site is paywalled, we may be able to use archive.is to fetch it."
|
||||
|
||||
- **What exists:** `recover: true` (`internal/recover`) already tries Wayback, then archive.today's
|
||||
`/newest/<url>` on archive.ph/.md/.li/.is, with fake-success detection. But it only fires on
|
||||
an *error* (`recoverWorthy`: hint `browser`, 422, 429, 504).
|
||||
- **The gap:** a soft paywall is usually a **200** with a teaser, so recover is never tried, and
|
||||
the agent gets a truncated article with no warning.
|
||||
- **Sketch:**
|
||||
1. Detect a paywall on a successful fetch. Signals: JSON-LD `isAccessibleForFree: false`,
|
||||
`hasPart.cssSelector` paywall markup, a known-paywall host list, and a clean body far shorter
|
||||
than the article's word count or the `og:` description implies.
|
||||
2. Add a `paywalled` warning and hint.
|
||||
3. With `recover`, try the snapshot routes. Mark the envelope `source.kind: "snapshot"`, as today.
|
||||
- **To measure first:** how often archive.today serves orion at all (it CAPTCHAs datacenter and
|
||||
tool traffic). Wayback excludes many paywalled publishers. The existing `/newest/` lookup
|
||||
only finds pages someone else has already archived; sift must not *submit* pages to
|
||||
archive.today.
|
||||
|
||||
### DataDome 401/403 — Reuters / Libération (titan, 2026-10-02)
|
||||
|
||||
Live examples: `https://www.reuters.com/business/what-is-frances-2027-budget-2026-10-01/`
|
||||
and `https://www.liberation.fr/international/amerique/face-a-liran-donald-trump-resserre-letau-militaire-economique-et-rhetorique-20261002_XX22YMZTQRE65BOXUUKD2ZKY3Y/`.
|
||||
|
||||
- **What sift returns:** Reuters `upstream 401` (no `hint`). Libération `upstream 403
|
||||
(likely bot-blocked); hint: browser`. P1 maps 403/429/451 (and 503-with-challenge) to
|
||||
bot-blocked; **401 is a generic `upstream %d`** with empty hint. Reuters AMP and Fusion
|
||||
(`/pf/api/v3/content/fetch/...`) 401 the same way — not an HTML-only wall.
|
||||
- **`--recover`:** both URLs have empty Wayback `available`. Reuters archive.today
|
||||
`/newest/` is 404, so recover rethrows the original 401 with no extra warning.
|
||||
Libération hits archive.ph `JlK7z` (snapshot 2026-10-02T16:40:46Z) and clean returns
|
||||
only `Please enable JS and disable any ad blocker` (`#cmsg`). That string is the
|
||||
publisher/DataDome stub; `fakeSuccess` does not reject it. PLAN's recover ladder still
|
||||
names an API/RSS/sitemap hop; `internal/recover` only does Wayback then archive.today.
|
||||
- **The wall is DataDome, not User-Agent.** From titan (`99.73.161.224`) Reuters is 401 and
|
||||
Libération 403 with `sift/682b9b8`, curl's default UA, and Chrome 131. Bodies are the
|
||||
same `var dd={…}` challenge (~770 bytes). Hermes browser on that IP sticks on
|
||||
DataDome (`🐴 reuters.com`, ~1.5 KB, empty article). No curl-impersonate/uTLS on titan
|
||||
to try a Chrome TLS fingerprint. A whole-service `user_agent` override or an automatic
|
||||
Mozilla retry would not have fetched either page from this egress.
|
||||
- **What did fetch:** Hermes `web_extract` (Firecrawl) returned the full Reuters explainer.
|
||||
Libération still paywalls after JS (lede + « réservé aux abonnés »). `sift search`
|
||||
Tavily snippets cover related Reuters copy, not that explainer URL.
|
||||
- **Open:** whether sift grows a DataDome-capable hop (Firecrawl as recover, browser or
|
||||
TLS-impersonate client on orion); whether 401 should be classified like 403; whether
|
||||
recover-miss and the `#cmsg` stub should be visible failures. Do not pick those here.
|
||||
Do not submit URLs to archive.today.
|
||||
|
||||
### Batch fetch endpoint (Kyle, 2026-09-26; plan after P5c)
|
||||
|
||||
A deliberate, **slow**, asynchronous endpoint for an agent that has collected a series of
|
||||
targets and wants them fetched without holding a request open per target. It builds on the
|
||||
batching infrastructure that will exist by the time P5c is deployed.
|
||||
|
||||
- **Submit:** `POST /v1/batch` with a list of targets (URLs, and presumably the same per-target
|
||||
options as `/v1/fetch`: `mode`, `question`, `pages`, `grep`, `jq`). Answers **202** with a job
|
||||
ID when accepted, **400** for a malformed request, **422** when a target is refused up front
|
||||
(for example the guard blocks it, or the list is over the cap).
|
||||
- **Poll and collect:** the caller checks in on the job ID later (`GET /v1/batch/{id}`: status and
|
||||
per-target progress) and fetches the results when they are done.
|
||||
- Questions for the design, when it is planned:
|
||||
- caps: targets per job, jobs per caller, total bytes;
|
||||
- per-target results and errors (one bad URL does not fail the job);
|
||||
- how long results are kept, and where (the SQLite jobs table the P5 pool discussion
|
||||
anticipated, under `/var/lib/sift`);
|
||||
- scheduling behind interactive requests (the same slots and gates, lower priority);
|
||||
- restricted callers: their jobs, results and logs partitioned like everything else;
|
||||
- cancellation and idempotency (resubmitting the same list);
|
||||
- whether Hermes needs a tool or skill change to use it.
|
||||
|
||||
@@ -11,7 +11,22 @@ One row per task session that ends in a commit. Newest last.
|
||||
| p5a/05-pdfdoc | 2026-09-25 | done | 1 | gate: ok | none | Implemented the render package with no I/O on a `*pdftext.Doc`. `Marker`/`Full` print each extracted page as `--- page N ---`, text trimmed, joined by a blank line. `Render` dispatches grep → page range → whole. Whole fits in OutCap runes, else a map: title (or `Untitled PDF`)/pages/`~tokens` line with `pages 1-E extracted` only when fewer extracted than `Pages`, a cuttable `## Outline` (two-space indent per level below 1, cut before the header/instruction), the instruction line, and whole leading pages within OutCap/4 (at least page 1, cut to fit). Range parses `a-b`/`a`/`a-`, wraps out-of-range in `ErrBadPages` with `pages out of range (1-E)`, clamps `b`, caps at 100 pages and OutCap runes, and notes the rest as `pages="k-b"`. Grep is RE2 with smart-case `(?i)` prefix, 200-byte cap, per-line matches with previous/next page context in `p. P-`/`p. P:` form, first 100 shown, total counted; RE2 keeps the `(a+)+$` test linear. Every length is in runes. Given test byte-identical; make gate ok. |
|
||||
| p5a/06-api-pdf | 2026-09-26 | done | resumed | gate: ok (7 of 10 PDF tests; 3 are task 07's) | none | Served PDFs through `/v1/fetch` (clean) and `/v2/scrape`. Added `RetryAfter int` to `fetch.Error` and a `Retry-After` header line guarded by `> 0`. Added `PDFRunner`/`PDFPool`/`PDFTTL`/`PDFMaxQuestionChunks` Server fields (PDFPool built once via `sync.Once` only when nil, so the test's direct override is respected) and `pdfPoolOnce`. Routed `clean.Kind` in `doFetch`/`cleanFetched` (html/json/text switched, "other"→415, "pdf"→extractPDF); `doFetch`'s switch routes "text" explicitly so a PDF never runs `clean.Text` first. New `internal/api/pdf.go`: `DefaultPDFTTL`/`DefaultPDFMaxQuestionChunks` (task 07), `runner()`/`pool()` defaults, `extractPDF` (pool key `callerClass|SHA256`, two-value assert, `fmt.Errorf("pdf: bad result %T")`), `pdfError` (422 for the five pdftext errors + panic default, 503+RetryAfter 10 for ErrBusy, 429 for ErrCallerBusy), `renderPDF` (`pdfdoc.Render`, whole text or map). `doFetch` returns the 4th `doc *pdftext.Doc`; digest uses `pdfdoc.Full(doc)` (computed once) as content, clean uses `renderPDF.Content`, envelope `pdf` object has pages/title/outline/extracted/first/last/map/grep_hits/note, raw mode→400. Scrape markdown=`renderPDF.Content`, html="", `metadata.numPages`=doc.Pages. 7 of 10 given tests pass under -race -count=3; the 3 failures (`TestPDFPagesAndGrep`, `TestPDFCache`, `TestPDFQuestionTooLong`) are task 07 tests that reference fields not in this task, exactly as the task doc (lines 145, 157-158) predicts; nothing else fails. Given files byte-identical. |
|
||||
| p5a/08-digest-cli-docs | 2026-09-26 | done | 1 | gate: ok | none | Added the page-cite sentence to the digest `systemPrompt` ("When the document has `--- page N ---` markers, cite the page of each fact as (p. N).") and set `PromptVersion = "3"`. Added the `sift_pdf_extract_total{outcome}`, `sift_pdf_extract_seconds`, `sift_pdf_pages_total`, `sift_pdf_pool_wait_seconds`, `sift_pdf_cache_hits_total` series to `internal/metrics` (`pdf.go` package) and wired `ObservePDF`/`ObservePDFWait`/`PDFCacheHit` into `internal/api/pdf.go` (`ObservePDF` once per call after `extractPDF`, outcome classified by `pdfOutcome` so raw text never becomes a label; `ObservePDFWait` only when the caller actually waited on `pool().Do`; `PDFCacheHit` on fresh hits and 304s). Refactored `cmd/sift/fetchcmd.go` into a `fetchOpts` struct with `parseFetchOpts` (returns `fetchOpts`, not `string, error`) and a `fetchBody(o fetchOpts) map[string]any` that does the digest/clean default and adds `pages`/`grep` only when set; `--pages`/`--grep` added to `fetchValueFlags`. Updated `skill/SKILL.md` (PDFs section; removed PDF mention from the 415 row). Unit: `nix/module.nix` `MemoryMax = "1G"` (one line). README: added a "PDF documents" section covering detection, limits (100 MB/120 s download, 60 s extraction, 2,000 pages), the child process, the 2-slot/4-wait/20 s pool, the 24 h conditional-GET cache, the answers, the task-06 errors table, and the metrics; noted the `go-pdfkit/extract` v0.3.0 engine chosen by the `docs/P5a.md` §3 bake-off. All tests pass under -race, make gate ok, three given files byte-identical, nix diff is one line. |
|
||||
| p5b/01-splitpages-bm25 | 2026-09-26 | done | 1 | gate: ok | none | `splitpages.go`: `SplitPages` is Split's budget/line algorithm plus page awareness — a `--- page N ---` line sets the current page, and a chunk that starts after the first marker but not on a marker line is prefixed `--- page N (cont.) ---\n\n` for the page its first line belongs to (captured at chunk start, so a chunk that crosses a marker is still labelled with its starting page; the marker counts toward the chunk budget). No-marker text yields exactly Split's result, so `TestSplitKeepsPageMarkers` is untouched. `bm25.go`: `BM25` with k1=1.2/b=0.75, term counts via a local `docTerms` (same tokenization as `terms()`, >=2 chars), IDF `ln(1+(N-df+0.5)/(df+0.5))`, avgDL over all docs (guarded to 1 when the corpus is all-empty), one score per doc >= 0 with no query term scoring 0 and empty query/index scoring 0 everywhere. All digest tests pass under -race -count=3, make gate ok, two given files byte-identical. |
|
||||
| p5b/02-retrieval | 2026-09-26 | done | 1 | gate: ok | none | Cached chunk embeddings + BM25-first retrieval. `digest.go`: `EmbedSyncMax=48`, `MaxDocChunks=5000`; `Result` gains `Retrieval string` and `Embedded int`; `selectText` now returns `(body, used, total, mode, embedded, err)` and `Run` sets both; `PutDigest` skips storing any `mode=="bm25"` answer (provisional, not cached). `retrieval.go`: `hybridOrBM25` — no cache, `len(chunks)<=48` embeds question+chunks inline and scores hybrid; `>48` scores BM25 and queues; cache present embeds the question and only the uncached chunks inline and scores hybrid on the union; both embed paths return hybrid, both bm25 paths return `bm25`. `background.go`: `bgWorker` with `bgOnce sync.Once` (creates the worker, never starts a goroutine) + `w.started` under `w.mu` (starts the goroutine once), dedup by `contentHash|callerClass`, full-queue drops the doc without marking seen so a later question re-queues, `acquireGate` retries once after 1 s on refusal, `WaitBackground` polls every 20 ms for true idle. `retrieval_test.go` copied byte-identical from the plan. All 12 digest tests pass under `-race -count=3`; `make gate` ok. |
|
||||
| p5a/07-api-cache-pages | 2026-09-26 | done | 1 | gate: ok | none | Implemented the PDF document cache with revalidation in `doFetch`: `pdfFromCache` (miss / fresh within `PDFTTL` / stale → `FetchIf` with 304→`TouchPDF`+answer-fresh / 200→return response with ok=false and no second GET) and `storePDF` after `extractPDF` (`ContentHash=res.SHA256`, ETag/Last-Modified from headers, FetchedAt=now, Meta=JSON of `doc.Meta`, Pages=`doc.Text`). `pages`/`grep` passed into `pdfdoc.Options` in `renderPDF`; `ErrBadPages`/`ErrBadGrep`→400 with the error's own text, non-PDF→400 "pages and grep apply to PDFs only" before the doc block. Question limit in digest mode: `n=len(digest.Split(pdfdoc.Full(doc), ChunkTokens))`, `n>PDFMaxQuestionChunks`→422 "document too long for question (n chunks); use grep or pages" before any digest call. Envelope `cached` set true for cache answers via the synthetic cache res's empty body (`len(res.Body)==0`); the one `doFetch` change beyond the rule-1 call was reusing that 200 response instead of a second `Fetch`. Added `fmt` import to api.go. Given file byte-identical; make gate ok. |
|
||||
| p5b/03-api | 2026-09-26 | done | 1 | gate: ok | none | The digest now answers long documents itself (task 02), so the API's refusal limit is only a safety cap: `DefaultPDFMaxQuestionChunks` 200→5000 in `internal/api/pdf.go` (the 422 message text unchanged, still "document too long for question (N chunks); use grep or pages"). In `v1fetch` (`internal/api/api.go`) the digest-meta map gains `"retrieval": dr.Retrieval` and `"embedded": dr.Embedded` next to `chunks_used`/`chunks_total`, only when the digest ran without a fallback. Copied the two given tests byte-identical from the plan — the replacement `pdf_test.go` is now the protected one and asserts `DefaultPDFMaxQuestionChunks == 5000`; `pdf_retrieval_test.go` asserts the first long-doc question returns `retrieval=bm25`, `embedded=0`. README PDF section: a long-doc question is answered first from BM25 (`retrieval: "bm25"`) while its chunks embed in the background and hybrid on later questions, plus the 5000-chunk cap. One sentence added to `skill/SKILL.md`. All api tests pass under `-race -count=3`; make gate ok; given files byte-identical. |
|
||||
| p5c/01-extract-api | 2026-09-26 | done | 1 | gate: ok | none | `POST /v1/extract` answers a PDF posted in the body like `/v1/fetch` answers the same PDF by URL, nothing fetched. `internal/api/extract.go`: auth via `caller`; `io.ReadAll` of the body (a `*http.MaxBytesError`→413 "too large (>N MB)", empty→400); `clean.Kind` must be "pdf" else 415; query `pages`/`grep`/`question`/`mode` (digest when a question is set, else clean; other→400). Content-addressed cache: `hash="sha256:"+HashBody(body)`, `key=PDFKey("upload:"+hash, callerClass)`; a `GetPDF` hit rebuilds the doc via `pdfFromEntry` with `cached:true` + `PDFCacheHit` and no child (bypasses the TTL/revalidation that `pdfFromCache` does, so identical content is extracted once); a miss builds a synthetic `fetch.Resource{URL:"upload",...}` and runs `extractPDF`+`storePDF`. Bad pages/grep→400 via `writeFetchErr`, the chunk limit→the same 422 as `v1fetch`, extraction errors→`pdfError` via `writeFetchErr`. Envelope mirrors the PDF fetch path (`content`/`mode`/`untrusted`/`warnings`/`sizes`/`timings`/`pdf`/`digest`); `source` is `{"kind":"upload","hash","bytes","cached","fetched_at"}` with no `url` key; no Honcho context for uploads; one log line "extract" with caller/bytes/pdf_pages/mode/cached. `internal/api/api.go` scoped to the rule-1 changes: `DefaultPDFMaxUpload` const (100 MiB), `PDFMaxUpload int64` Server field, the `POST /v1/extract` route in `New`, and `ServeHTTP` (per-path upload cap instead of the 64 KiB JSON limit, and `/v1/extract` added to the `requestDeadline` switch). The `pdfMaxUpload()` field-or-default resolver lives in extract.go. Given test byte-identical; all api tests pass under `-race -count=3`; make gate ok. |
|
||||
| p5c/02-extract-cli | 2026-09-26 | done | 1 | gate: ok | none | `sift extract <file.pdf>` uploads a local PDF to `/v1/extract` and prints the answer exactly as `sift fetch` does. `cmd/sift/extractcmd.go`: `extractOpts` + `parseExtractOpts` (reuses `splitAroundPositional` with `extractValueFlags` mode/pages/grep/question, one positional file, `--json`), `extractURL` builds `base+"/v1/extract"` with `url.Values` set only for the options that are set and `mode=digest` when a question is given and no mode — values are URL-encoded so a grep like `a b&c` survives round-trip; `extractMain` reuses `loadClient`, `stat`s the file and refuses >100 MiB before uploading (exit 3, matching the server's 413→exitFor), then `http.NewRequest` with the `*os.File` as the streaming body and `ContentLength` set (no second in-memory copy), `Content-Type: application/pdf` + bearer, and the same error/exitFor path as fetch; success prints `content` (or the `--json` envelope). `main.go`: `case "extract"` and one help line. README: extract line in the CLI block plus a sentence; SKILL.md: a PDF-on-disk bullet in the PDFs section. Given test byte-identical; all cmd/sift tests pass under `-race -count=3`; make gate ok; `bin/sift extract` prints the usage line. |
|
||||
| p6/01-batchstore | 2026-09-26 | done | 1 | gate: ok | none | Storage-only batch store in `internal/dcache` (new `batch.go`; schema added to `Open`; `BatchH`/`BatchM` counters added to `DB`). `CreateBatchJob` stores as given. Reads never write the lease: `BatchJob`, `BatchTargets` and the counts are pure, so exact lease values hold after a read. `BatchJobByIDem` returns the newest `(caller,idemKey)` job with `lease+ttl>now` and status queued/running/done/partial; failed/cancelled never reused. `PutBatchTargets` inserts in order; `BatchTargets` returns idx order with NULLs as nil pointers. `UpdateBatchTarget` bumps lease to `now`, increments `targets_done` exactly once per target entering done/failed (previous-status check prevents re-write double-count), and adds `len(result)` to `nbytes` when a result is written — all in one tx. `SetBatchJobStatus`/`TouchBatchJob` set lease; `SweepBatch` deletes by `lease+ttl<=now` with its targets and returns the job count; `EvictBatchOldest` evicts whole non-live jobs oldest-lease-first (never a queued/running one) until `BatchNbytes()<=maxBytes`; `MarkInFlightFailed` fails queued/running on restart; `BatchNbytes` sums job nbytes, kept out of `Size()`. Accessors return `(v,bool)`/`error`, count misses in `BatchM`, reuse `ClassFleet`/`ClassRestricted`. A first pass failed `TestBatchIDemLookup` because I forgot `.Scan(&id)` on the `QueryRow` (the un-scanned `*sql.Row` was always non-nil, so the method always missed); fixed and verified. All dcache tests pass under `-race -count=3`; make gate ok; given file byte-identical.
|
||||
| p6/02-runner-core | 2026-09-26 | done | 1 | gate: ok | caller-cleanup hand-off to task 03 | Runner core in new internal/batch with no workers (task 03 adds them; jobs stay queued, so a queued job counts as live for the caps). Public surface in batch.go: Options/Target/Result/TargetError, FetchFunc, Job/TargetStatus, Hooks and Config with defaults (TTL 24h, 32 targets/job, 1 job/caller, 8 total, worker clamped to 2, 64 MiB, 2 s/10 m retry). Submit runs the given order — empty→ErrNoTargets, idempotency via hex SHA-256 of `name+NUL+json.Marshal(targets)` (a live hit returns the id reused, touching no caps), then TargetsPerJob/JobsPerCaller/JobsTotal caps, then stores the job queued (random 16-byte hex id, CreatedAt=Lease=now) and its targets (pending, Options=json of t.Options) before returning so a queued job is live; the submitting caller is held in memory by job id, never persisted. Status maps rows and decodes Result without writing; Touch moves the lease; Cancel sets a queued job cancelled at once and leaves its targets pending; Close is idempotent. Store errors are Warn-logged with the job id, never a URL/host. |
|
||||
| p6/03-runner-workers | 2026-09-26 | done | 1 | gate: ok | none | Workers in new internal/batch/worker.go: one rotation of live jobs in submit order, round-robin via an absolute cursor (advances in pick, wraps at access, so it survives a second job being added while the first is in flight — the mod-n cursor from the design stuck at 0 for one job and re-picked it), serial per job (target i+1 never starts before i is stored), one target per fetch under the Close-cancelled context (panic→ErrInternal, never echoed; ErrRetryLater retried after RetryWait, bounded by RetryFor then stored failed with the wrapped text). Terminal status once per job: all done→done, all failed→failed, mix→partial. Cancel of a running job sets cancelled, stores the in-flight target without overwriting cancelled, leaves later targets pending. Close sets closed, cancels the context, broadcasts, waits the workers, idempotent. batch.go: New starts cfg.Workers, Submit appends the liveJob and wakes, Cancel handles the in-flight case, Close stops. Absolute cursor is the one deviation from the design's "index in live" comment. All race tests pass -count=3. | |
|
||||
| p6/04-metrics | 2026-09-26 | done | 1 | gate: ok | none | Batch series in new internal/metrics/batch.go, registered in New: `BatchJobsTotal`/`BatchTargetsTotal` CounterVecs (label "outcome", any value outside the known set folded to "other" via `batchJobOutcomes`/`batchTargetOutcomes` maps so a raw string never becomes a label), `BatchQueueDepth`/`BatchTargetsInflight` gauges, and `BatchJobSeconds`/`BatchTargetSeconds` histograms (durationBuckets) plus `BatchResultBytes` (1K,4K,16K,64K,256K,1M). Nil-receiver-safe `Observe…`/`Set…` methods mirroring ObservePDF/ObserveDigest. Series served through the existing /metrics handler; no new endpoint, no Snapshot change. Given test byte-identical; metrics tests pass under -race -count=3; make gate ok. |
|
||||
| p6/05-fetchtarget | 2026-09-26 | done | 1 | gate: ok | none | Extracted v1fetch's per-target pipeline into new `internal/api/target.go`: `fetchTarget(ctx, c, t batch.Target) (batch.Result, error)`. Body moved unchanged in order — question truncation to 500 runes, mode default/check, wantCtx, doFetch, recover path, PDF branch (render, pages/grep errors, question-chunk limit), clean/raw branch, digest, context search, the fetch log line, and the envelope. `content` is returned as `Result.Content`; every other envelope key (`mode`,`untrusted`,`source`,`warnings`,`sizes`,`timings`, and `pdf`/`digest`/`context` when present) goes in `Result.Meta`, never content. `v1fetch` keeps reqN/caller/decode/400, builds a `batch.Target` (`URL` + `Mode`/`Question`/`Pages`/`Grep`/`JQ`/`Context`/`Recover` into `Options`), calls `fetchTarget`, `writeFetchErr` on error and `writeJSON(w,200,env)` on success where `env` is a fresh copy of `Meta` plus `"content"` (a copy, so `len(env)==len(Meta)+1` holds and Meta is not mutated). The two `writeErr` failures become `*fetch.Error` with the same status/msg/hint: 400 "invalid url" (bad mode), 400 "raw mode does not apply to PDFs". `doFetch` takes `context.Context` instead of `*http.Request` (its three `r.Context()` → `ctx`); `scrape` passes `r.Context()`. Removed now-unused `fmt`/`unicode/utf8`/`recover` imports, added `batch`. All api tests pass under -race -count=3; make gate ok; given test byte-identical. |
|
||||
| p6/06-batch-wiring | 2026-09-26 | done | 1 | gate: ok | none | New `internal/api/batch.go`: `NewBatchRunner` overwrites `cfg.Hooks` so Job→ObserveBatchJob, Target→ObserveBatchTarget(batchOutcome)+ObserveBatchTargetSeconds (+ObserveBatchResultBytes only when err==nil), and JobSeconds/Depth/Inflight→the matching `s.Prom` methods; the runner's FetchFunc is `batchFetch`, which tags ctx with an unexported `batchCtxKey`, runs under `requestDeadline`, and maps fetchTarget's errors (a *fetch.Error 429/503 → `fmt.Errorf("%s: %w", fe.Msg, batch.ErrRetryLater)` so the runner retries; any other *fetch.Error → TargetError with 502 rewritten to 500 "internal error" and `s.blockN.Add(1)` for 403; anything else → TargetError 500 logged Error with `s.logErr`). `batchOutcome` classifies nil→ok, ErrRetryLater-wrapped→busy, then by TargetError.Status (403 blocked, 429 busy, 504 timeout, 415 unsupported, 422 whose Msg starts "document too long"→too_long, other 4xx→upstream_4xx, 5xx→upstream_5xx, any other error→upstream_5xx). Added `Batch *batch.Runner` Server field (nil → the /v1/batch endpoints answer 503 "batch disabled"). `pdf.go` extractPDF charges the PDF pool to `c.Name+"/batch"` when the ctx carries the batch mark, single-flight key unchanged so a batch and an interactive extract of the same PDF still share one child; the only new log line carries caller and host="-" (never a URL/host). All api tests pass under -race -count=3; make gate ok; given file byte-identical. |
|
||||
| p6/07-api-endpoints | 2026-09-26 | done | 1 | gate: ok | none | The four /v1/batch handlers in new `internal/api/batch.go`: `batchSubmit` (POST) decodes `{"targets":[{url,mode,question,pages,grep,jq,context,recover}]}` flat like `/v1/fetch`, 400 "invalid batch" on bad json/empty list/empty url/bad mode (only empty,clean,raw,digest), then `PreCheck` per target → 422 "blocked: target <idx>: <reason>" with no job, then `Submit` mapping ErrTooManyTargets→422 "too many targets (max 32)", ErrCallerBusy→429+Retry-After 30, ErrBusy→429+Retry-After 60, ErrNoTargets→400, and 202 {success,job_id,status,targets,reused} (status "queued" new, current status reused). `batchStatus` (GET) and `batchResult` (GET {id}/{idx}) and `batchCancel` (DELETE) share `ownership` (unknown or not `c.Name` → 404, never 403; nil batch → 503); status is a summary with no results unless `?results=1` inlines the envelope (Meta+content) on done targets, idx validated 0..targets_total-1 else 404, cancel of a terminal job → 409 "job already finished" else 200 {success,status:cancelled}. `api.go` scoped to the rule-2 changes: `batchMaxBody` (256 KiB) const, the four routes in `New`, and `ServeHTTP` applying the per-path cap once (`/v1/batch`→256 KiB, `/v1/extract`→pdfMaxUpload, else 64 KiB) and collapsing every batch path to a fixed metrics label via `batchObservePath` (`/v1/batch` or `/v1/batch/{id}`, never a job id). Each handler logs one `batch` line with caller/job_id/targets/status-or-reused and never a URL/hostname/idem_key. Given test byte-identical; all api tests pass under -race -count=3; make gate ok. |
|
||||
| p6/08-serve-config | 2026-09-26 | done | 1 | gate: ok | none | `[batch]` config in `internal/config`: the `Batch` type (disabled/targets_per_job/jobs_per_caller/jobs_total/workers/ttl/max_bytes) and `File.Batch`, plus `ParseSize` (trims spaces, accepts only binary KiB/MiB/GiB suffixes or a plain byte count, zero/negative/other → def). serve.go opens the dcache store when batch is on (`!cfg.Batch.Disabled`) alongside dixie/search — same path, same call, Warn on failure — then, before building the runner, reaps in-flight jobs a restart killed via `MarkInFlightFailed("restart", now)` with one `ObserveBatchJob("restart")` per reaped job and a `batch reaped` info log when n>0; builds `srv.Batch` from the config (TTL via `ParseDuration`, MaxBytes via `ParseSize`, runner defaults fill the zeros) only when enabled and the store opened; closes `srv.Batch` in both shutdown branches when non-nil. Added a commented `[batch]` block to dist/sift.toml.example. Given test byte-identical; all config and cmd tests pass under -race -count=3. |
|
||||
| p6/09-cli-docs | 2026-09-26 | done | 1 | gate: ok | none | `cmd/sift/batchcmd.go`: `parseBatchArgs` (verb + unchanged `rest`; submit takes one arg, status one job id optionally preceded by `--results`, get a job id plus a non-negative-integer idx, cancel one job id, else a usage error), `batchMain` dispatch, one small function per verb, and `batchDo` (bearer header, decode the JSON envelope, print it indented on stdout, on non-2xx print `env["error"]` to stderr and `exitFor`). `submitMain` reads the file or stdin and `buildTargetsBody` accepts a JSON object with a non-empty `targets` array, a bare array of target objects, or a bare array of URL strings (each wrapped `{"url": s}`), validating locally before the call. `main.go`: one `case "batch"` and one help line. README "Batch" section (four endpoints, serial-per-job/one-worker/round-robin model, 32 targets/1-per-caller/8-total caps, 256 KiB body, idempotency, lease/TTL/restart, summary-vs-results, restricted redaction, `sift_batch_*` metrics); SKILL.md "Batch" section. Given test byte-identical; all cmd/sift tests pass under -race -count=3. |
|
||||
| p3.1 | 2026-09-29 | done | 1 | gofmt/vet/race ok; make gate blocked by pre-existing staticcheck 0.6.1 vs go 1.26 | none | Native filters + multi-query RRF: Filter on Source.Search, SourceOut.Notes, Tavily auto_parameters false, SearXNG site:/time_range/categories, Brave notes, domain post-filter, Tavily min_score, sequential queries, JSON cache key, CLI flags, skill recipe. §11 tests in p31_*.go and search_filter_test.go. Guard/auth/orion pin untouched. |
|
||||
|
||||
## P5a review (owner, 2026-09-26)
|
||||
|
||||
|
||||
@@ -234,11 +234,14 @@ func TestGrepLimits(t *testing.T) {
|
||||
if _, err := pdfdoc.Render(d, pdfdoc.Options{Grep: "(unclosed", OutCap: 1000}); !errors.Is(err, pdfdoc.ErrBadGrep) {
|
||||
t.Errorf("bad regexp: %v, want ErrBadGrep", err)
|
||||
}
|
||||
// RE2 is linear: the classic catastrophic pattern over a 1 MB page finishes quickly.
|
||||
// RE2 is linear: the classic catastrophic pattern over a 1 MB page finishes. A backtracking
|
||||
// engine would not finish at all on this input, so the bound only has to be finite and far
|
||||
// above normal: 10 s. (1 s was too tight under -race, which slows regexp ~15–20x: talos
|
||||
// measured 1.5–1.9 s, 2026-09-26.)
|
||||
evil := &pdftext.Doc{Meta: pdftext.Meta{Pages: 1}, Text: []string{strings.Repeat("a", 1<<20) + "!"}}
|
||||
start := time.Now()
|
||||
_, _ = pdfdoc.Render(evil, pdfdoc.Options{Grep: "(a+)+$", OutCap: 1000})
|
||||
if d := time.Since(start); d > time.Second {
|
||||
if d := time.Since(start); d > 10*time.Second {
|
||||
t.Errorf("(a+)+$ took %v", d)
|
||||
}
|
||||
// Output stays under the cap.
|
||||
|
||||
@@ -0,0 +1,63 @@
|
||||
# P5b task 01: page-aware chunks and BM25
|
||||
|
||||
**Branch:** `p5b` (`git switch -c p5b master` if it does not exist, else `git switch p5b`; `git status --short` must be empty, otherwise stop)
|
||||
**Commit subject:** `digest: SplitPages keeps every chunk on its page; BM25 scorer`
|
||||
|
||||
## Goal
|
||||
|
||||
Two pure functions for task 02. `SplitPages` is `Split` for text with `--- page N ---`
|
||||
markers (what `pdfdoc.Full` produces): a chunk that starts inside a page begins with
|
||||
`--- page N (cont.) ---`, so every chunk names its page. `BM25` ranks chunks lexically.
|
||||
|
||||
## Files
|
||||
|
||||
- Copy: `internal/digest/splitpages_test.go`, `internal/digest/bm25_test.go`
|
||||
- Create: `internal/digest/splitpages.go`, `internal/digest/bm25.go`
|
||||
- Modify: `docs/implementer-log.md` only
|
||||
|
||||
## Interfaces
|
||||
|
||||
```go
|
||||
package digest
|
||||
|
||||
// SplitPages splits text into chunks of at most maxTok tokens (EstTokens) on line boundaries,
|
||||
// like Split. A marker line matching ^--- page (\d+) ---$ sets the current page; a chunk that
|
||||
// starts after the first marker but not at a marker line begins with
|
||||
// "--- page N (cont.) ---\n\n" for the current page N. Text with no marker lines gives exactly
|
||||
// Split's result.
|
||||
func SplitPages(text string, maxTok int) []string
|
||||
|
||||
type BM25 struct{ /* unexported */ }
|
||||
|
||||
// NewBM25 indexes docs (k1 = 1.2, b = 0.75). Terms are the same tokens as terms() — lower-case
|
||||
// runs of letters and digits, at least 2 characters — but counted, not as a set.
|
||||
func NewBM25(docs []string) *BM25
|
||||
|
||||
// Scores returns one score per document, in order, each >= 0 and never NaN or Inf. An empty
|
||||
// query, an empty index or a document without any query term scores 0.
|
||||
func (b *BM25) Scores(query string) []float64
|
||||
```
|
||||
|
||||
## Rules
|
||||
|
||||
1. `SplitPages`: the continuation marker counts toward the chunk's budget (the test allows a
|
||||
marker line of slack). Never split inside a line. Do not change `Split` (a P5a given test
|
||||
pins it).
|
||||
2. BM25 IDF: `ln(1 + (N − df + 0.5) / (df + 0.5))`, which is always positive. Average document
|
||||
length over all documents; an all-empty corpus scores 0 everywhere.
|
||||
|
||||
## Steps
|
||||
|
||||
- [ ] **1.** Branch; copy the two given tests. **2. See them fail** (undefined names).
|
||||
- [ ] **3.** `splitpages.go`, **4.** `bm25.go`.
|
||||
- [ ] **5.** `gofmt -w`; `go test -race -count=3 ./internal/digest/` → `ok`. **6.** `make gate`.
|
||||
- [ ] **7.** Row `p5b/01-splitpages-bm25`; commit.
|
||||
|
||||
```sh
|
||||
git add internal/digest docs/implementer-log.md
|
||||
git commit -F .state/commit-msg.txt
|
||||
```
|
||||
|
||||
## Done when
|
||||
|
||||
- The given tests pass, every existing digest test passes, `make gate` ok, given files byte-identical.
|
||||
@@ -0,0 +1,105 @@
|
||||
# P5b task 02: cached embeddings and BM25-first retrieval for long documents
|
||||
|
||||
**Branch:** `p5b` (`git switch p5b`; `git status --short` must be empty, otherwise stop)
|
||||
**Commit subject:** `digest: cache chunk embeddings; long documents answer from BM25 while embedding in the background`
|
||||
|
||||
## Goal
|
||||
|
||||
Today `selectText` embeds every chunk of a document for every question, inline, inside the
|
||||
digest's 45 s timeout, and never uses the embedding cache. After this task:
|
||||
- chunk embeddings are cached;
|
||||
- a document with few uncached chunks is embedded inline and scored **hybrid** (as today);
|
||||
- a long one is answered at once from **BM25**, with **no embedding call on the request path**,
|
||||
while its chunks are embedded by a background worker;
|
||||
- once they are, the next question is hybrid.
|
||||
|
||||
## Files
|
||||
|
||||
- Copy: `internal/digest/retrieval_test.go`
|
||||
- Create: `internal/digest/retrieval.go` (selection), `internal/digest/background.go` (the worker)
|
||||
- Modify: `internal/digest/digest.go` (`Result` fields, `Run`, `selectText`'s call site),
|
||||
`docs/implementer-log.md`
|
||||
|
||||
## Interfaces
|
||||
|
||||
```go
|
||||
package digest
|
||||
|
||||
const (
|
||||
EmbedSyncMax = 48 // uncached chunks a question may embed inline; more go to the background
|
||||
MaxDocChunks = 5000 // background embedding cap per document
|
||||
)
|
||||
|
||||
// in Result:
|
||||
Retrieval string // "whole", "head", "hybrid", "bm25" or "cached" (see rule 2)
|
||||
Embedded int // chunks that had an embedding when the answer was built
|
||||
|
||||
// WaitBackground returns when the background worker is idle with nothing queued, or ctx.Err().
|
||||
func (e *Engine) WaitBackground(ctx context.Context) error
|
||||
```
|
||||
|
||||
What already exists and stays:
|
||||
- `Split`, `SplitPages`, `NewBM25` (task 01);
|
||||
- `EstTokens`, `SmallDoc`, `ChunkTokens`, `EmbedBatch`, `lexicalOverlap`, `cosine`;
|
||||
- `e.embed(ctx, inputs)` (batches of `EmbedBatch`, halves on "too large");
|
||||
- `e.Gate.Acquire/Release`;
|
||||
- `dcache.EmbedKey(contentHash string, chunkIndex int, embedModel, callerClass string)`, `e.Cache.GetEmbed(key)` / `PutEmbed(key, vec)` (a nil `*DB` is a no-op).
|
||||
|
||||
## Rules
|
||||
|
||||
1. **Chunks and keys.**
|
||||
- Chunk with `SplitPages(cleaned, ChunkTokens)`.
|
||||
- Embedding keys are `dcache.EmbedKey(contentHash+"/sp1200", i, e.embedModel(), callerClass)`: the suffix binds the key to this chunker, so a future chunking change cannot reuse wrong vectors.
|
||||
- `contentHash` and `callerClass` come from `Run` (pass them to the selection).
|
||||
2. **Selection** (`retrieval.go`), in order:
|
||||
- `EstTokens(cleaned) <= SmallDoc` → the whole text, `Retrieval "whole"`.
|
||||
- Empty question → the leading chunks within `SmallDoc`, as today, `Retrieval "head"`.
|
||||
- **No cache** (`e.Cache == nil`):
|
||||
- `len(chunks) <= EmbedSyncMax` → embed everything inline, as today, `"hybrid"`;
|
||||
- otherwise `"bm25"`, with no embedding and nothing queued (there is nowhere to keep the vectors).
|
||||
- **Cache:** look up every chunk's vector.
|
||||
- **uncached `<= EmbedSyncMax`** → embed the question and the uncached chunks inline (the question first, then the missing chunks, through `e.embed`), store each new chunk vector, `"hybrid"`.
|
||||
- **otherwise** → `"bm25"`: rank by `NewBM25(chunks).Scores(question)` normalised by its maximum. Make **no embedding call**, and queue the document for the background (rule 4).
|
||||
- Hybrid score = `0.7·cosine(q, v_i) + 0.3·lexicalOverlap(question, chunk)`, as today.
|
||||
- Both modes then use today's greedy budget: chunk 0 first, then best-scoring chunks while the total stays within `SmallDoc`, emitted in document order.
|
||||
- `Embedded` = the chunks with a vector at that moment.
|
||||
3. **`Run`:**
|
||||
- Set `r.Retrieval` and `r.Embedded`.
|
||||
- On a digest-cache hit: `Retrieval "cached"`.
|
||||
- **A digest made in `"bm25"` mode is not stored in the digest cache.** It is provisional; once the embeddings exist the same question must get a hybrid answer.
|
||||
4. **Background worker** (`background.go`):
|
||||
- One goroutine per `Engine`, started on first use (`sync.Once`), with a queue of at most **4 documents**, deduplicated by `contentHash+"|"+callerClass`. A full queue drops the document, logs `Warn` (with no text) and a later question re-queues it.
|
||||
- For each document: embed the first `MaxDocChunks` chunks that lack a vector, **`EmbedBatch` at a time**. For each batch:
|
||||
- `e.Gate.Acquire(context.Background())`, so interactive digests get turns; on a refusal (queue full or wait), sleep 1 s and try again;
|
||||
- `e.embed` with a `e.timeout()` context;
|
||||
- `Release`;
|
||||
- store the vectors.
|
||||
- On an embedding error: log `Warn`, drop the rest of that document, go on to the next.
|
||||
- `WaitBackground` is true-idle: the queue is empty and no document is in progress. Poll every 20 ms or use a condition variable; either is fine.
|
||||
5. Never log document text. The worker never holds the gate between batches.
|
||||
|
||||
## Facts
|
||||
|
||||
- The given test's fake dixie counts every embedded input. It can **hold `/embed` shut**, which is how it proves the `"bm25"` answer made no embedding call on the request path: if yours does, the first `Run` blocks and the test fails on time.
|
||||
- `sseHandler` and `chatReq` come from `digest_test.go` / `digest.go` in the same package.
|
||||
- Existing tests (`TestChunkOrderAndBudget`, `TestEmbedTooLargeRetriesAndCompletes`, …) run with no cache and must stay green.
|
||||
|
||||
## Steps
|
||||
|
||||
- [ ] **1.** `git switch p5b`; copy the given test. **2. See it fail.**
|
||||
- [ ] **3.** `retrieval.go`, **4.** `background.go`, **5.** the `Run`/`Result` changes.
|
||||
- [ ] **6.** `gofmt -w`; `go test -race -count=3 ./internal/digest/` → `ok`. **7.** `make gate`.
|
||||
- [ ] **8.** Row `p5b/02-retrieval`; commit.
|
||||
|
||||
```sh
|
||||
git add internal/digest docs/implementer-log.md
|
||||
git commit -F .state/commit-msg.txt
|
||||
```
|
||||
|
||||
## Done when
|
||||
|
||||
- The given and existing digest tests pass under `-race -count=3`; `make gate` ok; given file byte-identical.
|
||||
|
||||
## Stop and report if
|
||||
|
||||
- A given test's timing cannot be met without sleeping on the request path.
|
||||
@@ -0,0 +1,43 @@
|
||||
# P5b task 03: the API answers long-document questions and says how
|
||||
|
||||
**Branch:** `p5b` (`git switch p5b`; `git status --short` must be empty, otherwise stop)
|
||||
**Commit subject:** `api: questions on long PDFs are answered (retrieval in the digest meta); README and skill`
|
||||
|
||||
## Goal
|
||||
|
||||
With task 02, the digest handles long documents itself. The API's 200-chunk refusal becomes a
|
||||
5000-chunk safety cap (task 02's `MaxDocChunks`), and the `/v1/fetch` envelope's `digest`
|
||||
object says how the answer was built.
|
||||
|
||||
## Files
|
||||
|
||||
- Copy: `internal/api/pdf_retrieval_test.go`, and the **replacement** `internal/api/pdf_test.go`
|
||||
(P5a's copy asserted `DefaultPDFMaxQuestionChunks == 200`; now 5000 — the only change; the
|
||||
P5b copy is the protected one from now on)
|
||||
- Modify: `internal/api/pdf.go` (the constant), `internal/api/api.go` (the two digest-meta keys
|
||||
only), `README.md`, `skill/SKILL.md`, `docs/implementer-log.md`
|
||||
|
||||
## Rules
|
||||
|
||||
1. `DefaultPDFMaxQuestionChunks = 5000`. The 422 message is unchanged.
|
||||
2. In `v1fetch`, when the digest ran without a fallback, `digMeta` gains
|
||||
`"retrieval": dr.Retrieval` and `"embedded": dr.Embedded`, next to `chunks_used` and
|
||||
`chunks_total`.
|
||||
3. **`README.md`:** in the PDF section, a question on a long document is answered first from BM25 (`retrieval: "bm25"`) while its chunks are embedded in the background, and is hybrid on later questions. Also mention the 5000-chunk cap.
|
||||
4. **`skill/SKILL.md`:** one sentence: `--question` works on long PDFs; the first answer may be `bm25` (lexical); ask again later for a sharper one.
|
||||
|
||||
## Steps
|
||||
|
||||
- [ ] **1.** `git switch p5b`; copy the two given files. **2. See them fail.**
|
||||
- [ ] **3.** The constant, the two meta keys, README, skill.
|
||||
- [ ] **4.** `gofmt -w`; `go test -race -count=3 ./internal/api/` → `ok`. **5.** `make gate`.
|
||||
- [ ] **6.** Row `p5b/03-api`; commit.
|
||||
|
||||
```sh
|
||||
git add internal/api README.md skill/SKILL.md docs/implementer-log.md
|
||||
git commit -F .state/commit-msg.txt
|
||||
```
|
||||
|
||||
## Done when
|
||||
|
||||
- All tests pass; `make gate` ok; given files byte-identical.
|
||||
@@ -0,0 +1,35 @@
|
||||
# P5b implementation plan: questions on long documents
|
||||
|
||||
> **For the implementing model:** do not work from this file. The owner gives you one task file
|
||||
> at a time. This file is the index for the owner and the reviewer.
|
||||
|
||||
**Goal:** P5a refuses `question` on a PDF longer than 200 digest chunks (~240k tokens) with
|
||||
"use grep or pages". After P5b such a question is answered: at once from **BM25** over
|
||||
page-aware chunks, while the document's chunks are **embedded in the background**; once they
|
||||
are, later questions use today's hybrid score. Two facts found while planning:
|
||||
|
||||
- The embedding cache (`dcache` `GetEmbed`/`PutEmbed`, `EmbedKey`) exists but **nothing uses it**:
|
||||
every question re-embeds every chunk.
|
||||
- `Split` is line-based, so a chunk that starts mid-page carries no `--- page N ---` marker and
|
||||
the model cannot cite its page.
|
||||
|
||||
| Task | Package | Given |
|
||||
|---|---|---|
|
||||
| 01-splitpages-bm25 | `internal/digest` (pure functions) | `digest/splitpages_test.go`, `digest/bm25_test.go` |
|
||||
| 02-retrieval | `internal/digest` | `digest/retrieval_test.go` |
|
||||
| 03-api | `internal/api`, `README.md`, `skill/SKILL.md` | `api/pdf_retrieval_test.go`; replaces `api/pdf_test.go` (P5a: the 200-chunk default becomes 5000) |
|
||||
|
||||
**How this plan was made:** acceptance tests first, no reference implementation; the given tests
|
||||
compiled against a panic-only skeleton on master `1311cf3` and failed for the intended reasons.
|
||||
|
||||
## Changes during the run
|
||||
|
||||
- **Run 1, task 01 (2026-09-26 11:01–12:24): infrastructure, not the model.** straylight's
|
||||
llama-server ran out of shared KV (4 slots over one 262,144-token pool). The session got
|
||||
`Context size has been exceeded` with no commit. Hermes's ornith context was capped at 131072
|
||||
(Kyle: "b is acceptable"); the task was re-run clean and passed (`dca0004`, rebased to `7595ad8`).
|
||||
- **Task 02, first attempt: my test, not the model or the task.** `TestEmbeddingsAreCachedAcrossQuestions`
|
||||
used `longDoc(12)`, which is 4,617 tokens. That is under `SmallDoc` (6000), so rule 2 answers
|
||||
it `"whole"`, while the test expected `"hybrid"`. Ornith measured the document and was working
|
||||
it out when I stopped the session (22 min). Fix `d52a4fa`: `longDoc(30)` (11,677 tokens, 10
|
||||
chunks) plus a `t.Fatalf` precondition. Relaunched from task 02 on a clean tree.
|
||||
@@ -0,0 +1,75 @@
|
||||
package api
|
||||
|
||||
// P5b task 03: a question on a long PDF is answered (by BM25 first) instead of refused, and the
|
||||
// envelope says how the digest was built.
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/digest"
|
||||
"git.wntrmute.dev/kyle/sift/internal/pdftest"
|
||||
)
|
||||
|
||||
// fakeModel answers /embed with a fixed vector and /chat with a one-word SSE reply.
|
||||
func fakeModel(t *testing.T) *httptest.Server {
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/embed", func(w http.ResponseWriter, r *http.Request) {
|
||||
var body struct {
|
||||
Input []string `json:"input"`
|
||||
}
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
var data []map[string]any
|
||||
for i := range body.Input {
|
||||
data = append(data, map[string]any{"index": i, "embedding": []float64{1, float64(i % 7)}})
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"data": data})
|
||||
})
|
||||
mux.HandleFunc("/chat", func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = w.Write([]byte(`data: {"choices":[{"delta":{"content":"answer (p. 3)"}}]}` + "\n\ndata: [DONE]\n\n"))
|
||||
})
|
||||
srv := httptest.NewServer(mux)
|
||||
t.Cleanup(srv.Close)
|
||||
return srv
|
||||
}
|
||||
|
||||
func TestPDFLongQuestionIsAnswered(t *testing.T) {
|
||||
r := newPDFRig(t)
|
||||
m := fakeModel(t)
|
||||
db := r.s.Digest.Cache
|
||||
r.s.Digest = &digest.Engine{ChatURL: m.URL + "/chat", EmbedURL: m.URL + "/embed", Cache: db, Gate: digest.NewGate(2, 8, 5*time.Second), Client: m.Client()}
|
||||
var pages []string
|
||||
for p := 1; p <= 260; p++ { // ~260 pages × 40 lines: well over EmbedSyncMax chunks
|
||||
pages = append(pages, pdftest.Lorem(p, 40))
|
||||
}
|
||||
o := newOrigin(t, "application/pdf", pdftest.Build(pdftest.Doc{Title: "Long", Pages: pages}))
|
||||
code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL, "question": "what is on page three?", "context": false})
|
||||
if code != 200 {
|
||||
t.Fatalf("long question: %d %v, want 200 (P5a refused this)", code, out)
|
||||
}
|
||||
dg, _ := out["digest"].(map[string]any)
|
||||
if dg == nil || dg["retrieval"] != "bm25" || dg["embedded"] != 0.0 {
|
||||
t.Errorf("digest meta = %v, want retrieval bm25 and embedded 0 on the first question", dg)
|
||||
}
|
||||
if c, _ := out["content"].(string); !strings.Contains(c, "answer") {
|
||||
t.Errorf("content %q", c)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFSmallQuestionReportsRetrieval(t *testing.T) {
|
||||
r := newPDFRig(t)
|
||||
m := fakeModel(t)
|
||||
db := r.s.Digest.Cache
|
||||
r.s.Digest = &digest.Engine{ChatURL: m.URL + "/chat", EmbedURL: m.URL + "/embed", Cache: db, Gate: digest.NewGate(2, 8, 5*time.Second), Client: m.Client()}
|
||||
o := newOrigin(t, "application/pdf", threePDF())
|
||||
code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL, "question": "what is on page two?", "context": false})
|
||||
dg, _ := out["digest"].(map[string]any)
|
||||
if code != 200 || dg == nil || dg["retrieval"] != "whole" {
|
||||
t.Errorf("small document: %d digest %v, want retrieval whole", code, dg)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,403 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
package digest
|
||||
|
||||
// P5b task 01: BM25, the lexical ranking used while a long document's embeddings are still
|
||||
// being built (and inside the hybrid score after).
|
||||
|
||||
import (
|
||||
"math"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBM25RanksTheRareTerm(t *testing.T) {
|
||||
docs := []string{
|
||||
"the ship sailed from helsinki to tallinn in the morning",
|
||||
"the ship stayed in port all day",
|
||||
"a zanzibarite sample was found near the harbour",
|
||||
"the harbour master logged the ship",
|
||||
}
|
||||
s := NewBM25(docs).Scores("where was the zanzibarite found")
|
||||
if len(s) != len(docs) {
|
||||
t.Fatalf("%d scores for %d docs", len(s), len(docs))
|
||||
}
|
||||
best := 0
|
||||
for i := range s {
|
||||
if s[i] > s[best] {
|
||||
best = i
|
||||
}
|
||||
}
|
||||
if best != 2 {
|
||||
t.Errorf("best = doc %d (%v), want doc 2, the only one with the rare term", best, s)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBM25Properties(t *testing.T) {
|
||||
filler := strings.Repeat("word ", 50)
|
||||
docs := []string{
|
||||
"alpha " + filler, // 0: tf 1
|
||||
strings.Repeat("alpha ", 10) + filler, // 1: tf 10, same filler
|
||||
"alpha common", // 2: tf 1, short
|
||||
"common " + filler, // 3: no alpha
|
||||
}
|
||||
idx := NewBM25(docs)
|
||||
s := idx.Scores("alpha")
|
||||
if !(s[1] > s[0]) {
|
||||
t.Errorf("more occurrences must score higher: tf10 %v vs tf1 %v", s[1], s[0])
|
||||
}
|
||||
if !(s[1] < 10*s[0]) {
|
||||
t.Errorf("term frequency must saturate: tf10 %v, 10 × tf1 %v", s[1], 10*s[0])
|
||||
}
|
||||
if !(s[2] > s[0]) {
|
||||
t.Errorf("the same term in a shorter document must score higher: short %v, long %v", s[2], s[0])
|
||||
}
|
||||
if s[3] != 0 {
|
||||
t.Errorf("a document without the term scores %v, want 0", s[3])
|
||||
}
|
||||
// IDF: a term in every document is worth less than a rare one.
|
||||
all := NewBM25([]string{"xx yy", "xx zz", "xx ww"}).Scores("xx")
|
||||
rare := NewBM25([]string{"xx yy", "qq zz", "qq ww"}).Scores("xx")
|
||||
if !(rare[0] > all[0]) {
|
||||
t.Errorf("idf: a term in 1 of 3 docs must outscore one in 3 of 3: %v vs %v", rare[0], all[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestBM25Degenerate(t *testing.T) {
|
||||
for _, docs := range [][]string{nil, {}, {""}, {"", "a b"}} {
|
||||
s := NewBM25(docs).Scores("a")
|
||||
if len(s) != len(docs) {
|
||||
t.Errorf("%q: %d scores", docs, len(s))
|
||||
}
|
||||
for _, v := range s {
|
||||
if math.IsNaN(v) || math.IsInf(v, 0) || v < 0 {
|
||||
t.Errorf("%q: score %v", docs, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, v := range NewBM25([]string{"a b", "c d"}).Scores("") {
|
||||
if v != 0 {
|
||||
t.Errorf("empty query scored %v", v)
|
||||
}
|
||||
}
|
||||
// Deterministic.
|
||||
a := NewBM25([]string{"alpha beta", "beta gamma", "gamma alpha"}).Scores("alpha gamma")
|
||||
b := NewBM25([]string{"alpha beta", "beta gamma", "gamma alpha"}).Scores("alpha gamma")
|
||||
for i := range a {
|
||||
if a[i] != b[i] {
|
||||
t.Fatal("scores differ between identical runs")
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,247 @@
|
||||
package digest
|
||||
|
||||
// P5b task 02: retrieval for long documents. Embeddings are cached per chunk; a document with
|
||||
// more uncached chunks than EmbedSyncMax is answered at once from BM25 while its chunks are
|
||||
// embedded in the background, and later questions use the hybrid score.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"hash/fnv"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
)
|
||||
|
||||
// fakeDixie serves /embed (a bag-of-words hash into 64 dims, so cosine similarity means
|
||||
// something) and /chat (the SSE stub). It counts embedded inputs and can hold /embed shut.
|
||||
type fakeDixie struct {
|
||||
srv *httptest.Server
|
||||
embedded atomic.Int64 // inputs embedded, questions included
|
||||
hold chan struct{}
|
||||
holdOnce sync.Once
|
||||
lastChat atomic.Pointer[chatReq]
|
||||
}
|
||||
|
||||
func newFakeDixie(t *testing.T, hold bool) *fakeDixie {
|
||||
f := &fakeDixie{}
|
||||
if hold {
|
||||
f.hold = make(chan struct{})
|
||||
}
|
||||
var disc atomic.Bool
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/embed", func(w http.ResponseWriter, r *http.Request) {
|
||||
if f.hold != nil {
|
||||
select {
|
||||
case <-f.hold:
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
var body struct {
|
||||
Input []string `json:"input"`
|
||||
}
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
var data []map[string]any
|
||||
for i, s := range body.Input {
|
||||
data = append(data, map[string]any{"index": i, "embedding": bow(s)})
|
||||
}
|
||||
f.embedded.Add(int64(len(body.Input)))
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"data": data})
|
||||
})
|
||||
mux.Handle("/chat", sseHandler(t, "ok", &f.lastChat, &disc))
|
||||
f.srv = httptest.NewServer(mux)
|
||||
t.Cleanup(func() { f.release(); f.srv.Close() })
|
||||
return f
|
||||
}
|
||||
|
||||
func (f *fakeDixie) release() {
|
||||
if f.hold != nil {
|
||||
f.holdOnce.Do(func() { close(f.hold) })
|
||||
}
|
||||
}
|
||||
|
||||
func bow(s string) []float64 {
|
||||
v := make([]float64, 64)
|
||||
for t := range terms(s) {
|
||||
h := fnv.New32a()
|
||||
_, _ = h.Write([]byte(t))
|
||||
v[h.Sum32()%64]++
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func (f *fakeDixie) engine(t *testing.T, gate *Gate, withCache bool) *Engine {
|
||||
e := &Engine{ChatURL: f.srv.URL + "/chat", EmbedURL: f.srv.URL + "/embed", Gate: gate, Client: f.srv.Client()}
|
||||
if withCache {
|
||||
db, err := dcache.Open(filepath.Join(t.TempDir(), "c.db"), 0, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
e.Cache = db
|
||||
}
|
||||
return e
|
||||
}
|
||||
|
||||
// longDoc is n pages of distinct filler; the page named rare carries a rare term.
|
||||
func longDoc(n, rare int) string {
|
||||
var b strings.Builder
|
||||
words := []string{"harbour", "vessel", "cargo", "anchor", "pilot", "tug", "berth", "quay", "hull", "keel", "mast", "deck"}
|
||||
for p := 1; p <= n; p++ {
|
||||
if p > 1 {
|
||||
b.WriteString("\n\n")
|
||||
}
|
||||
fmt.Fprintf(&b, "--- page %d ---\n\n", p)
|
||||
for l := 0; l < 40; l++ {
|
||||
fmt.Fprintf(&b, "page %d line %d %s %s %s %s\n", p, l, words[(p+l)%12], words[(p*3+l)%12], words[(p*5+l)%12], words[(p*7+l)%12])
|
||||
}
|
||||
if p == rare {
|
||||
b.WriteString("the zanzibarite cargo manifest was signed by the second officer\n")
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func chatDoc(f *fakeDixie) string {
|
||||
req := f.lastChat.Load()
|
||||
if req == nil {
|
||||
return ""
|
||||
}
|
||||
var s strings.Builder
|
||||
for _, m := range req.Messages {
|
||||
s.WriteString(m.Content)
|
||||
}
|
||||
return s.String()
|
||||
}
|
||||
|
||||
func TestEmbeddingsAreCachedAcrossQuestions(t *testing.T) {
|
||||
f := newFakeDixie(t, false)
|
||||
e := f.engine(t, NewGate(2, 8, 5*time.Second), true)
|
||||
doc := longDoc(30, 0) // ~11.5k tokens: over SmallDoc (so not "whole"), ~10 chunks: under EmbedSyncMax
|
||||
if EstTokens(doc) <= SmallDoc {
|
||||
t.Fatalf("test document is %d tokens; it must exceed SmallDoc (%d) to reach retrieval", EstTokens(doc), SmallDoc)
|
||||
}
|
||||
r := e.Run(context.Background(), doc, "which berth did the tug use?", "https://ex.test/a.pdf", dcache.ClassFleet)
|
||||
if r.Fallback != "" || r.Retrieval != "hybrid" {
|
||||
t.Fatalf("first question: fallback %q retrieval %q", r.Fallback, r.Retrieval)
|
||||
}
|
||||
n := len(SplitPages(doc, ChunkTokens))
|
||||
if got := f.embedded.Load(); got != int64(n+1) {
|
||||
t.Fatalf("first question embedded %d inputs, want %d chunks + the question", got, n)
|
||||
}
|
||||
r = e.Run(context.Background(), doc, "where was the pilot?", "https://ex.test/a.pdf", dcache.ClassFleet)
|
||||
if r.Fallback != "" || r.Retrieval != "hybrid" || r.Embedded != n {
|
||||
t.Fatalf("second question: %+v", r)
|
||||
}
|
||||
if got := f.embedded.Load(); got != int64(n+2) {
|
||||
t.Errorf("second question embedded %d inputs in total, want %d: only its own question is new", got, n+2)
|
||||
}
|
||||
}
|
||||
|
||||
// A long document is answered from BM25 without waiting for a single embedding, and becomes
|
||||
// hybrid once the background has embedded every chunk.
|
||||
func TestLongDocumentBM25ThenHybrid(t *testing.T) {
|
||||
f := newFakeDixie(t, true) // /embed held shut until release
|
||||
e := f.engine(t, NewGate(2, 8, 5*time.Second), true)
|
||||
doc := longDoc(300, 250)
|
||||
n := len(SplitPages(doc, ChunkTokens))
|
||||
if n <= EmbedSyncMax {
|
||||
t.Fatalf("test document has %d chunks; it must exceed EmbedSyncMax (%d)", n, EmbedSyncMax)
|
||||
}
|
||||
start := time.Now()
|
||||
r := e.Run(context.Background(), doc, "who signed the zanzibarite manifest?", "https://ex.test/big.pdf", dcache.ClassFleet)
|
||||
if d := time.Since(start); d > 10*time.Second {
|
||||
t.Fatalf("first answer took %v: it waited on embeddings", d)
|
||||
}
|
||||
if r.Fallback != "" || r.Retrieval != "bm25" || r.Embedded != 0 {
|
||||
t.Fatalf("first answer: %+v, want retrieval bm25 with 0 embedded", r)
|
||||
}
|
||||
if doc := chatDoc(f); !strings.Contains(doc, "zanzibarite") || !strings.Contains(doc, "page 250") {
|
||||
t.Errorf("BM25 did not bring page 250 (the only one with the term) into the prompt")
|
||||
}
|
||||
|
||||
f.release()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
||||
defer cancel()
|
||||
if err := e.WaitBackground(ctx); err != nil {
|
||||
t.Fatalf("background embedding did not finish: %v", err)
|
||||
}
|
||||
r = e.Run(context.Background(), doc, "who signed the zanzibarite manifest?", "https://ex.test/big.pdf", dcache.ClassFleet)
|
||||
if r.Fallback != "" || r.Retrieval != "hybrid" || r.Embedded != n {
|
||||
t.Fatalf("after background: %+v, want hybrid with all %d embedded", r, n)
|
||||
}
|
||||
if doc := chatDoc(f); !strings.Contains(doc, "zanzibarite") {
|
||||
t.Errorf("hybrid lost the page with the term")
|
||||
}
|
||||
}
|
||||
|
||||
// Two questions on the same long document queue it once: every chunk is embedded once.
|
||||
func TestBackgroundDedupes(t *testing.T) {
|
||||
f := newFakeDixie(t, true)
|
||||
e := f.engine(t, NewGate(2, 8, 5*time.Second), true)
|
||||
doc := longDoc(200, 0)
|
||||
n := len(SplitPages(doc, ChunkTokens))
|
||||
e.Run(context.Background(), doc, "first question about the quay", "https://ex.test/b.pdf", dcache.ClassFleet)
|
||||
e.Run(context.Background(), doc, "second question about the keel", "https://ex.test/b.pdf", dcache.ClassFleet)
|
||||
f.release()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
||||
defer cancel()
|
||||
if err := e.WaitBackground(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := f.embedded.Load(); got > int64(n) {
|
||||
t.Errorf("embedded %d inputs for a %d-chunk document: queued twice", got, n)
|
||||
}
|
||||
}
|
||||
|
||||
// Background work takes the gate one batch at a time, so an interactive digest is not starved.
|
||||
func TestBackgroundYieldsTheGate(t *testing.T) {
|
||||
f := newFakeDixie(t, false)
|
||||
e := f.engine(t, NewGate(1, 8, 10*time.Second), true)
|
||||
e.Run(context.Background(), longDoc(400, 0), "start the background", "https://ex.test/c.pdf", dcache.ClassFleet)
|
||||
start := time.Now()
|
||||
r := e.Run(context.Background(), "a short page about a tug and a quay", "what is here?", "https://ex.test/small", dcache.ClassFleet)
|
||||
if r.Fallback != "" {
|
||||
t.Fatalf("interactive digest fell back (%q) while background embedding ran", r.Fallback)
|
||||
}
|
||||
if d := time.Since(start); d > 5*time.Second {
|
||||
t.Errorf("interactive digest waited %v behind background embedding", d)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
||||
defer cancel()
|
||||
_ = e.WaitBackground(ctx)
|
||||
}
|
||||
|
||||
// Without a cache there is nowhere to keep background work: a long document is BM25 and
|
||||
// nothing is embedded.
|
||||
func TestNoCacheLongDocumentIsBM25Only(t *testing.T) {
|
||||
f := newFakeDixie(t, false)
|
||||
e := f.engine(t, NewGate(2, 8, 5*time.Second), false)
|
||||
r := e.Run(context.Background(), longDoc(300, 250), "who signed the zanzibarite manifest?", "https://ex.test/d.pdf", dcache.ClassFleet)
|
||||
if r.Fallback != "" || r.Retrieval != "bm25" {
|
||||
t.Fatalf("%+v, want bm25", r)
|
||||
}
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
if got := f.embedded.Load(); got != 0 {
|
||||
t.Errorf("embedded %d inputs with no cache to keep them in", got)
|
||||
}
|
||||
}
|
||||
|
||||
// A document that fits the digest budget whole is sent whole, as today.
|
||||
func TestSmallDocumentIsWhole(t *testing.T) {
|
||||
f := newFakeDixie(t, false)
|
||||
e := f.engine(t, NewGate(2, 8, 5*time.Second), true)
|
||||
r := e.Run(context.Background(), "--- page 1 ---\n\nOne short page.", "what is it?", "https://ex.test/e.pdf", dcache.ClassFleet)
|
||||
if r.Fallback != "" || r.Retrieval != "whole" {
|
||||
t.Fatalf("%+v, want whole", r)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
package digest
|
||||
|
||||
// P5b task 01: SplitPages keeps every chunk of a PDF's text attached to its page, so a digest
|
||||
// built from a chunk that starts mid-page can still cite the page.
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// pagesText builds n pages of about perPage tokens each, every line naming its page.
|
||||
func pagesText(n, perPage int) string {
|
||||
var b strings.Builder
|
||||
for p := 1; p <= n; p++ {
|
||||
if p > 1 {
|
||||
b.WriteString("\n\n")
|
||||
}
|
||||
fmt.Fprintf(&b, "--- page %d ---\n\n", p)
|
||||
for l := 0; l*12 < perPage; l++ {
|
||||
fmt.Fprintf(&b, "page %d line %d alpha bravo charlie delta echo foxtrot golf hotel\n", p, l)
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
var markerRE = regexp.MustCompile(`^--- page (\d+)( \(cont\.\))? ---`)
|
||||
|
||||
func TestSplitPagesEveryChunkNamesItsPage(t *testing.T) {
|
||||
text := pagesText(4, 3000) // each page ~3000 tokens: several chunks per page at 1200
|
||||
chunks := SplitPages(text, ChunkTokens)
|
||||
if len(chunks) < 8 {
|
||||
t.Fatalf("%d chunks, want at least 8 for 4 × 3000 tokens at %d", len(chunks), ChunkTokens)
|
||||
}
|
||||
cont := 0
|
||||
for i, c := range chunks {
|
||||
m := markerRE.FindStringSubmatch(c)
|
||||
if m == nil {
|
||||
t.Fatalf("chunk %d does not start with a page marker: %.80q", i, c)
|
||||
}
|
||||
if m[2] != "" {
|
||||
cont++
|
||||
// A continuation marker names the page the chunk's first line belongs to.
|
||||
first := strings.SplitN(strings.TrimPrefix(c, m[0]), "\n", 3)
|
||||
line := strings.TrimSpace(strings.Join(first, " "))
|
||||
if !strings.Contains(line, "page "+m[1]+" line ") {
|
||||
t.Errorf("chunk %d marked page %s but starts with %.60q", i, m[1], line)
|
||||
}
|
||||
}
|
||||
if EstTokens(c) > ChunkTokens+20 {
|
||||
t.Errorf("chunk %d is %d tokens, over the %d budget (plus a marker line)", i, EstTokens(c), ChunkTokens)
|
||||
}
|
||||
}
|
||||
if cont == 0 {
|
||||
t.Error("no continuation markers: pages longer than a chunk must produce some")
|
||||
}
|
||||
}
|
||||
|
||||
// Nothing is lost or duplicated: removing the continuation markers gives back the original lines.
|
||||
func TestSplitPagesPreservesContent(t *testing.T) {
|
||||
text := pagesText(3, 2500)
|
||||
var got []string
|
||||
for _, c := range SplitPages(text, ChunkTokens) {
|
||||
for _, line := range strings.Split(c, "\n") {
|
||||
if strings.HasSuffix(line, " (cont.) ---") || strings.TrimSpace(line) == "" {
|
||||
continue
|
||||
}
|
||||
got = append(got, line)
|
||||
}
|
||||
}
|
||||
var want []string
|
||||
for _, line := range strings.Split(text, "\n") {
|
||||
if strings.TrimSpace(line) != "" {
|
||||
want = append(want, line)
|
||||
}
|
||||
}
|
||||
if strings.Join(got, "\n") != strings.Join(want, "\n") {
|
||||
t.Errorf("content changed: %d lines back, %d in", len(got), len(want))
|
||||
}
|
||||
}
|
||||
|
||||
// Text with no markers splits exactly as Split does.
|
||||
func TestSplitPagesWithoutMarkersIsSplit(t *testing.T) {
|
||||
text := strings.Repeat("plain words without any page markers at all\n", 2000)
|
||||
a, b := SplitPages(text, ChunkTokens), Split(text, ChunkTokens)
|
||||
if strings.Join(a, "\x00") != strings.Join(b, "\x00") {
|
||||
t.Errorf("SplitPages differs from Split on unmarked text: %d vs %d chunks", len(a), len(b))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
# P5c task 01: `POST /v1/extract`
|
||||
|
||||
**Branch:** `p5c` (`git switch -c p5c master` if it does not exist, else `git switch p5c`; `git status --short` must be empty, otherwise stop)
|
||||
**Commit subject:** `api: POST /v1/extract answers an uploaded PDF like a fetched one (content-addressed cache)`
|
||||
|
||||
## Goal
|
||||
|
||||
`POST /v1/extract` with a PDF as the request body returns the `/v1/fetch` PDF envelope:
|
||||
text with page markers or the map, `pages`, `grep`, or a cited digest. Nothing is fetched.
|
||||
The bytes are hashed, the cache key is the hash, and the same file is extracted once.
|
||||
|
||||
## Files
|
||||
|
||||
- Copy: `internal/api/extract_test.go`
|
||||
- Create: `internal/api/extract.go` (the handler and its envelope)
|
||||
- Modify: `internal/api/api.go`: only the route in `New`, the `Server` field and constant
|
||||
below, and the two lines in `ServeHTTP` (rule 1); `docs/implementer-log.md`
|
||||
|
||||
## Interfaces
|
||||
|
||||
```go
|
||||
package api
|
||||
|
||||
const DefaultPDFMaxUpload = 100 << 20
|
||||
|
||||
// in Server:
|
||||
PDFMaxUpload int64 // 0 = DefaultPDFMaxUpload
|
||||
|
||||
// route: s.Mux.HandleFunc("POST /v1/extract", s.v1extract)
|
||||
// query parameters (all optional): pages, grep, question, mode ("clean" | "digest";
|
||||
// default digest when question is set, else clean)
|
||||
```
|
||||
|
||||
What exists and is reused, unchanged:
|
||||
- `s.caller`, `s.callerClass`;
|
||||
- `s.extractPDF(ctx, c, *fetch.Resource)`, `pdfError`, `s.renderPDF`, `s.storePDF(key, res, doc)`, `s.pdfFromEntry`;
|
||||
- `s.cache()`, `dcache.PDFKey`, `clean.Kind`, `pdfdoc.Full`;
|
||||
- `s.Digest.Run`, `s.Prom.PDFCacheHit`, `writeErr`, `writeJSON`.
|
||||
|
||||
## Rules
|
||||
|
||||
1. **`ServeHTTP`:**
|
||||
- today it wraps every request body in `http.MaxBytesReader(w, r.Body, maxReqBody)` (64 KiB). For `/v1/extract` use `PDFMaxUpload` instead;
|
||||
- add `/v1/extract` to the paths that get `requestDeadline`;
|
||||
- every other path keeps 64 KiB (the given test checks that).
|
||||
2. **Handler:**
|
||||
- Authenticate as the other endpoints do.
|
||||
- Read the whole body. Empty → `400 empty body`. Over the limit (`*http.MaxBytesError`) → `413 too large (>N MB)`.
|
||||
- `clean.Kind(Content-Type, body)` must be `"pdf"`; otherwise `415 extract accepts PDFs only`.
|
||||
3. **Cache:**
|
||||
- `hash := "sha256:" + hex(sha256(body))`, `key := dcache.PDFKey("upload:"+hash, s.callerClass(c))`.
|
||||
- An entry → the doc from it, `cached: true`, `PDFCacheHit`, no child. The content is addressed by its hash, so there is **no TTL and no revalidation**.
|
||||
- Otherwise build `&fetch.Resource{URL: "upload", FinalURL: "upload", Status: 200, ContentType: ct, Body: body, Hash: hash, FetchedAt: now}`, then `extractPDF`, then `storePDF(key, res, doc)`.
|
||||
- Errors go through `pdfError`, via the existing `writeFetchErr` (it adds `Retry-After`).
|
||||
4. **Answer:** the same fields as `/v1/fetch` for a PDF:
|
||||
- `content`, `mode`, `untrusted: true`, `warnings`, `sizes`, `timings`, `pdf`, and `digest` when a digest ran, including P5b's `retrieval` and `embedded`;
|
||||
- `source` is `{"kind": "upload", "hash", "bytes", "cached", "fetched_at"}`, with **no `url`** key;
|
||||
- bad `pages`/`grep` → 400 (as in `v1fetch`); the question limit → the same 422 as `v1fetch`.
|
||||
- Do not call Honcho context for uploads.
|
||||
5. **Logging:** one line, `"extract"`, with caller, bytes, `pdf_pages`, mode and cached. Never a title or any text. The server never sees a filename.
|
||||
|
||||
## Steps
|
||||
|
||||
- [ ] **1.** Branch; copy the given test. **2. See it fail.**
|
||||
- [ ] **3.** `api.go` hooks, **4.** `extract.go`.
|
||||
- [ ] **5.** `gofmt -w`; `go test -race -count=3 ./internal/api/` → `ok`. **6.** `make gate`.
|
||||
- [ ] **7.** Row `p5c/01-extract-api`; commit.
|
||||
|
||||
```sh
|
||||
git add internal/api docs/implementer-log.md
|
||||
git commit -F .state/commit-msg.txt
|
||||
```
|
||||
|
||||
## Done when
|
||||
|
||||
- All tests pass under `-race -count=3`; `make gate` ok; given file byte-identical.
|
||||
@@ -0,0 +1,60 @@
|
||||
# P5c task 02: `sift extract <file.pdf>`
|
||||
|
||||
**Branch:** `p5c` (`git switch p5c`; `git status --short` must be empty, otherwise stop)
|
||||
**Commit subject:** `cli: sift extract uploads a local PDF; README and skill`
|
||||
|
||||
## Goal
|
||||
|
||||
`sift extract thesis.pdf [--pages 41-60] [--grep RE] [--question Q] [--json]` uploads the file
|
||||
to `/v1/extract` and prints the answer exactly as `sift fetch` does.
|
||||
|
||||
## Files
|
||||
|
||||
- Copy: `cmd/sift/extract_test.go`
|
||||
- Create: `cmd/sift/extractcmd.go`
|
||||
- Modify: `cmd/sift/main.go` (a `case "extract":` and one help line), `README.md`,
|
||||
`skill/SKILL.md`, `docs/implementer-log.md`
|
||||
|
||||
## Interfaces
|
||||
|
||||
```go
|
||||
package main
|
||||
|
||||
type extractOpts struct {
|
||||
File, Pages, Grep, Question, Mode string
|
||||
JSON bool
|
||||
}
|
||||
|
||||
const extractUsage = "usage: sift extract <file.pdf> [--pages a-b] [--grep RE] [--question Q] [--mode clean|digest] [--json]"
|
||||
|
||||
func parseExtractOpts(args []string) (extractOpts, error) // flags before or after the file, as fetch
|
||||
func extractURL(base string, o extractOpts) string // base + "/v1/extract?" + only the options that are set; mode=digest when Question is set and Mode is empty
|
||||
```
|
||||
|
||||
## Rules
|
||||
|
||||
1. The command:
|
||||
- Resolve the base URL and key exactly as `sift fetch` does (reuse its client-config helpers; do not copy them).
|
||||
- `stat` the file: over 100 MiB → exit with a message before uploading.
|
||||
- `POST` the bytes with `Content-Type: application/pdf` and the bearer key.
|
||||
- Print `content` (or the whole envelope with `--json`), with the same errors and exit codes as fetch (`exitFor`).
|
||||
- Do not read the whole file into memory twice: stream it with `os.Open` as the request body and set `ContentLength`.
|
||||
2. **`README.md`:** a `sift extract` subsection next to `sift fetch`.
|
||||
3. **`skill/SKILL.md`:** "A PDF on disk (an attachment, a download): `sift extract <file>`, with the same `--pages`/`--grep`/`--question`. Never run `pdftotext` yourself."
|
||||
|
||||
## Steps
|
||||
|
||||
- [ ] **1.** `git switch p5c`; copy the given test. **2. See it fail.**
|
||||
- [ ] **3.** `extractcmd.go`, `main.go`, README, skill.
|
||||
- [ ] **4.** `gofmt -w`; `go test -race -count=3 ./cmd/sift/` → `ok`. **5.** `make gate`.
|
||||
- [ ] **6.** `make build && bin/sift extract 2>&1 | head -1` prints the usage line.
|
||||
- [ ] **7.** Row `p5c/02-extract-cli`; commit.
|
||||
|
||||
```sh
|
||||
git add cmd/sift README.md skill/SKILL.md docs/implementer-log.md
|
||||
git commit -F .state/commit-msg.txt
|
||||
```
|
||||
|
||||
## Done when
|
||||
|
||||
- All tests pass; `make gate` ok; given file byte-identical.
|
||||
@@ -0,0 +1,21 @@
|
||||
# P5c implementation plan: PDFs already on a host
|
||||
|
||||
> **For the implementing model:** do not work from this file. The owner gives you one task file
|
||||
> at a time. This file is the index for the owner and the reviewer.
|
||||
|
||||
**Goal:** an agent with a PDF on disk (a Mattermost attachment, a download) gets the same
|
||||
answer as for a PDF URL — text with page markers, the map, pages, grep, cited digests — by
|
||||
uploading the bytes: `POST /v1/extract`, CLI `sift extract <file>`. No poppler on the host, no
|
||||
shell `pdftotext`. The cache key is the content hash, so the same file is extracted once.
|
||||
|
||||
| Task | Package | Given |
|
||||
|---|---|---|
|
||||
| 01-extract-api | `internal/api` | `api/extract_test.go` |
|
||||
| 02-extract-cli | `cmd/sift`, `README.md`, `skill/SKILL.md` | `cmd/sift/extract_test.go` |
|
||||
|
||||
Runs after P5b (task 01 uses the P5b digest fields). Given tests compiled against a panic-only
|
||||
skeleton on master `1311cf3` and failed for the intended reasons.
|
||||
|
||||
## Changes during the run
|
||||
|
||||
(none yet)
|
||||
@@ -0,0 +1,46 @@
|
||||
package main
|
||||
|
||||
// P5c task 02: `sift extract <file>` uploads a PDF to /v1/extract.
|
||||
|
||||
import (
|
||||
"net/url"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestParseExtractOpts(t *testing.T) {
|
||||
o, err := parseExtractOpts([]string{"thesis.pdf", "--pages", "41-60", "--grep", "attention head"})
|
||||
if err != nil || o.File != "thesis.pdf" || o.Pages != "41-60" || o.Grep != "attention head" {
|
||||
t.Errorf("%+v %v", o, err)
|
||||
}
|
||||
o, err = parseExtractOpts([]string{"--question", "what is the method?", "--json", "paper.pdf"})
|
||||
if err != nil || o.File != "paper.pdf" || o.Question != "what is the method?" || !o.JSON {
|
||||
t.Errorf("flags first: %+v %v", o, err)
|
||||
}
|
||||
if _, err := parseExtractOpts(nil); err == nil {
|
||||
t.Error("no file: want an error")
|
||||
}
|
||||
if _, err := parseExtractOpts([]string{"a.pdf", "b.pdf"}); err == nil {
|
||||
t.Error("two files: want an error")
|
||||
}
|
||||
if !strings.Contains(extractUsage, "sift extract <file.pdf>") {
|
||||
t.Errorf("usage %q", extractUsage)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractURL(t *testing.T) {
|
||||
u, err := url.Parse(extractURL("https://sift.example:8443", extractOpts{File: "x.pdf", Pages: "2-3", Grep: "a b&c"}))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if u.Path != "/v1/extract" || u.Query().Get("pages") != "2-3" || u.Query().Get("grep") != "a b&c" {
|
||||
t.Errorf("url %s", u)
|
||||
}
|
||||
if u.Query().Has("question") || u.Query().Has("mode") {
|
||||
t.Errorf("unset options must not be sent: %s", u)
|
||||
}
|
||||
u, _ = url.Parse(extractURL("https://sift.example:8443", extractOpts{File: "x.pdf", Question: "why?"}))
|
||||
if u.Query().Get("question") != "why?" || u.Query().Get("mode") != "digest" {
|
||||
t.Errorf("question url %s (mode must be digest with a question)", u)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,122 @@
|
||||
package api
|
||||
|
||||
// P5c task 01: POST /v1/extract takes a PDF's bytes (a file already on the caller's host), and
|
||||
// answers exactly like /v1/fetch does for a PDF URL. The cache key is the content hash.
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func (r *pdfRig) extract(t *testing.T, key, query string, body []byte, ct string) (int, map[string]any) {
|
||||
t.Helper()
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/extract"+query, bytes.NewReader(body))
|
||||
if ct != "" {
|
||||
req.Header.Set("Content-Type", ct)
|
||||
}
|
||||
if key != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
r.s.ServeHTTP(rec, req)
|
||||
var out map[string]any
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &out)
|
||||
return rec.Code, out
|
||||
}
|
||||
|
||||
func TestExtractUpload(t *testing.T) {
|
||||
r := newPDFRig(t)
|
||||
pdf := threePDF()
|
||||
code, out := r.extract(t, r.fleet, "", pdf, "application/pdf")
|
||||
if code != 200 {
|
||||
t.Fatalf("%d %v", code, out)
|
||||
}
|
||||
content, _ := out["content"].(string)
|
||||
if !strings.Contains(content, "--- page 2 ---") || !strings.Contains(content, "quick brown fox") {
|
||||
t.Errorf("content %q", content)
|
||||
}
|
||||
if p := pdfInfo(t, out); p["pages"] != 3.0 || p["title"] != "Three Pages" {
|
||||
t.Errorf("pdf %v", p)
|
||||
}
|
||||
src, _ := out["source"].(map[string]any)
|
||||
if src["kind"] != "upload" || src["bytes"] != float64(len(pdf)) || !strings.HasPrefix(fmt.Sprint(src["hash"]), "sha256:") || src["cached"] != false {
|
||||
t.Errorf("source %v", src)
|
||||
}
|
||||
if _, hasURL := src["url"]; hasURL {
|
||||
t.Error("an upload has no URL; source must not invent one")
|
||||
}
|
||||
// The same bytes again: from the cache, no child.
|
||||
code, out = r.extract(t, r.fleet, "?pages=2", pdf, "application/octet-stream")
|
||||
if code != 200 || r.runs() != 1 {
|
||||
t.Fatalf("second upload: %d, %d extractions; want 200 and 1", code, r.runs())
|
||||
}
|
||||
if src, _ := out["source"].(map[string]any); src["cached"] != true {
|
||||
t.Errorf("second upload source %v, want cached", src)
|
||||
}
|
||||
if c, _ := out["content"].(string); strings.Contains(c, "--- page 1 ---") || !strings.Contains(c, "--- page 2 ---") {
|
||||
t.Errorf("pages=2: %q", c)
|
||||
}
|
||||
code, out = r.extract(t, r.fleet, "?grep=quick+brown", pdf, "application/pdf")
|
||||
if c, _ := out["content"].(string); code != 200 || !strings.Contains(c, "p. 2: The quick brown fox") {
|
||||
t.Errorf("grep: %d %q", code, c)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractRefusals(t *testing.T) {
|
||||
r := newPDFRig(t)
|
||||
if code, _ := r.extract(t, "", "", threePDF(), "application/pdf"); code != 401 {
|
||||
t.Errorf("no key: %d, want 401", code)
|
||||
}
|
||||
if code, out := r.extract(t, r.fleet, "", []byte("<html>not a pdf</html>"), "text/html"); code != 415 {
|
||||
t.Errorf("html body: %d %v, want 415", code, out)
|
||||
}
|
||||
if code, _ := r.extract(t, r.fleet, "", nil, "application/pdf"); code != 400 {
|
||||
t.Errorf("empty body: %d, want 400", code)
|
||||
}
|
||||
if code, _ := r.extract(t, r.fleet, "?pages=0-1", threePDF(), "application/pdf"); code != 400 {
|
||||
t.Errorf("bad pages: %d, want 400", code)
|
||||
}
|
||||
r.s.PDFMaxUpload = 1000
|
||||
if code, out := r.extract(t, r.fleet, "", threePDF(), "application/pdf"); code != 413 {
|
||||
t.Errorf("over PDFMaxUpload: %d %v, want 413", code, out)
|
||||
}
|
||||
if DefaultPDFMaxUpload != 100<<20 {
|
||||
t.Errorf("DefaultPDFMaxUpload = %d, want 100 MiB", DefaultPDFMaxUpload)
|
||||
}
|
||||
// The ordinary JSON endpoints keep their small request limit.
|
||||
big := bytes.Repeat([]byte("x"), 128<<10)
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/fetch", bytes.NewReader(append([]byte(`{"url":"`), append(big, []byte(`"}`)...)...)))
|
||||
req.Header.Set("Authorization", "Bearer "+r.fleet)
|
||||
r.s.ServeHTTP(rec, req)
|
||||
if rec.Code == 200 {
|
||||
t.Error("a 128 KiB /v1/fetch body was accepted; the JSON limit must still apply there")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractRestrictedPartition(t *testing.T) {
|
||||
r := newPDFRig(t)
|
||||
pdf := threePDF()
|
||||
if code, _ := r.extract(t, r.gs, "", pdf, "application/pdf"); code != 200 {
|
||||
t.Fatalf("restricted upload: %d", code)
|
||||
}
|
||||
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 })
|
||||
if strings.Contains(b.String(), "Three Pages") || strings.Contains(b.String(), "quick brown") {
|
||||
t.Errorf("an upload's log line carries document text: %q", b.String())
|
||||
}
|
||||
}
|
||||
r.logs.mu.Unlock()
|
||||
if code, _ := r.extract(t, r.fleet, "", pdf, "application/pdf"); code != 200 || r.runs() != 2 {
|
||||
t.Errorf("fleet after restricted: %d, %d extractions; want 200 and 2 (separate partitions)", code, r.runs())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,155 @@
|
||||
# P6 task 01: batch tables, lease, GC, byte budget in `internal/dcache`
|
||||
|
||||
**Branch:** `p6` (`git switch -c p6 master` if it does not exist, else `git switch p6`; `git status --short` must be empty, otherwise stop)
|
||||
**Commit subject:** `dcache: batch job and target tables with a sliding lease, sweep and byte budget`
|
||||
|
||||
Add the batch job store to the existing `internal/dcache` SQLite database. No HTTP, no worker —
|
||||
this task is storage only. `dcache` owns the single SQLite connection, so the batch tables and
|
||||
their accessors live here; the runner (task 02) and the API (tasks 06–07) call them.
|
||||
|
||||
## Context
|
||||
|
||||
`dcache.DB` already opens one SQLite connection (`SetMaxOpenConns(1)`, WAL) in `Open`, owns
|
||||
`maxBytes`, `evict()`, `Size()`, and `Stats()`. The batch store is a sibling, not a member of the
|
||||
content-addressed cache: a batch job is an atomic unit (a row plus its target rows) and must be
|
||||
evicted whole by lease, never row-by-row by the content LRU.
|
||||
|
||||
## Schema (add to `Open`)
|
||||
|
||||
```sql
|
||||
CREATE TABLE IF NOT EXISTS batch_job (
|
||||
id TEXT PRIMARY KEY,
|
||||
caller TEXT NOT NULL,
|
||||
caller_class TEXT NOT NULL, -- 'fleet' | 'restricted'
|
||||
idem_key TEXT NOT NULL,
|
||||
status TEXT NOT NULL, -- queued|running|done|partial|failed|cancelled
|
||||
error TEXT NOT NULL DEFAULT '',
|
||||
targets_total INTEGER NOT NULL,
|
||||
targets_done INTEGER NOT NULL DEFAULT 0,
|
||||
created_at INTEGER NOT NULL, -- unix seconds
|
||||
lease INTEGER NOT NULL, -- unix seconds; bumped on submit/completion/read
|
||||
nbytes INTEGER NOT NULL DEFAULT 0 -- sum of stored per-target result bytes
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS batch_target (
|
||||
job_id TEXT NOT NULL,
|
||||
idx INTEGER NOT NULL, -- order within the job, 0..n-1
|
||||
url TEXT NOT NULL,
|
||||
options TEXT NOT NULL, -- JSON: the per-target options
|
||||
status TEXT NOT NULL, -- pending|running|done|failed
|
||||
result TEXT, -- JSON result envelope, or NULL
|
||||
error TEXT, -- error message, or NULL
|
||||
hint TEXT, -- error hint, or NULL
|
||||
PRIMARY KEY (job_id, idx)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS batch_job_idem ON batch_job(caller, idem_key);
|
||||
CREATE INDEX IF NOT EXISTS batch_job_lease ON batch_job(lease);
|
||||
```
|
||||
|
||||
`result`/`error`/`hint` are `NULL` until the target finishes. The `result` JSON is the
|
||||
`/v1/fetch`-shaped envelope, bounded by the output cap at write time (the render step enforces the cap;
|
||||
the store just stores what it is given and counts its bytes into `nbytes`).
|
||||
|
||||
## Public surface (exact names — the given test compiles against them)
|
||||
|
||||
```go
|
||||
type BatchJob struct {
|
||||
ID string
|
||||
Caller string
|
||||
CallerClass string
|
||||
IDemKey string
|
||||
Status string
|
||||
Error string
|
||||
TargetsTotal int
|
||||
TargetsDone int
|
||||
CreatedAt int64
|
||||
Lease int64
|
||||
NBytes int64
|
||||
}
|
||||
|
||||
type BatchTarget struct {
|
||||
JobID string
|
||||
Idx int
|
||||
URL string
|
||||
Options string
|
||||
Status string
|
||||
Result *string
|
||||
Error *string
|
||||
Hint *string
|
||||
}
|
||||
|
||||
func (d *DB) CreateBatchJob(j BatchJob) error
|
||||
func (d *DB) BatchJob(id string) (BatchJob, bool)
|
||||
func (d *DB) BatchJobByIDem(caller, idemKey string, now int64, ttl time.Duration) (string, bool)
|
||||
func (d *DB) LiveBatchJobs(caller string) int
|
||||
func (d *DB) LiveBatchJobsTotal() int
|
||||
func (d *DB) PutBatchTargets(jobID string, targets []BatchTarget) error
|
||||
func (d *DB) BatchTargets(jobID string) ([]BatchTarget, bool)
|
||||
func (d *DB) UpdateBatchTarget(jobID string, idx int, now int64, status string, result, errMsg, hint *string) error
|
||||
func (d *DB) SetBatchJobStatus(id, status, errMsg string, now int64) error
|
||||
func (d *DB) TouchBatchJob(id string, now int64) error
|
||||
func (d *DB) SweepBatch(now int64, ttl time.Duration) int
|
||||
func (d *DB) EvictBatchOldest(maxBytes int64) int
|
||||
func (d *DB) MarkInFlightFailed(reason string, now int64) int
|
||||
func (d *DB) BatchNbytes() int64
|
||||
```
|
||||
|
||||
Every method that touches time takes an explicit `now` (unix seconds) — the store is pure and the
|
||||
caller decides the clock, so the given test asserts exact lease values.
|
||||
|
||||
## Rules
|
||||
|
||||
1. **Sliding lease.** `CreateBatchJob` stores the job as given (the caller sets `Lease`/`CreatedAt`).
|
||||
`UpdateBatchTarget`, `SetBatchJobStatus`, `TouchBatchJob`, and `MarkInFlightFailed` set the job
|
||||
lease to their `now`. **Reads never write:** `BatchJob`, `BatchTargets` and the counts change
|
||||
nothing, so the given test can assert exact lease values after reading. A consumer's read keeps
|
||||
a job alive because the API calls `TouchBatchJob` after it (task 07), not because the read does.
|
||||
2. **No partial-job reads leak a competing caller.** Every accessor is keyed by `job_id` only; the
|
||||
caller filter is the API layer's job (task 07). The store stores what it is told and returns
|
||||
what is asked.
|
||||
3. **Idempotency lookup.** `BatchJobByIDem(caller, idemKey, now, ttl)` returns the id of a job with
|
||||
that `(caller, idem_key)` whose `lease + ttl > now` **and whose status is `queued`, `running`,
|
||||
`done` or `partial`**. A `failed` job (including one killed by a restart) or a `cancelled` one
|
||||
is never returned: resubmitting the same list after either starts a new job. If several match,
|
||||
return the newest by `created_at`.
|
||||
4. `BatchTargets` returns targets in `idx` order. `UpdateBatchTarget` increments `targets_done`
|
||||
exactly once per target that transitions *into* `done`/`failed` (a re-write of the same target
|
||||
does not double-count), bumps the job lease to `now`, and adds `len(result)` to `nbytes` when a
|
||||
result is written.
|
||||
5. `SweepBatch(now, ttl)` deletes every job whose `lease + ttl <= now`, and its target rows, and
|
||||
returns the number of jobs deleted. It is the only thing that removes a finished job by time; the runner calls it (task 02).
|
||||
6. `EvictBatchOldest(maxBytes)` evicts **whole jobs** (oldest `lease` first) until `BatchNbytes()`
|
||||
is `<= maxBytes`, and returns the number evicted. A job in `queued`/`running` is never
|
||||
byte-evicted (only lease-GC, cancel, or restart removes an active job). The runner (task 02) calls it
|
||||
after each target result is written; it is separate from the content-cache `evict()`.
|
||||
7. `MarkInFlightFailed(reason, now)` marks every `queued`/`running` job `failed` with
|
||||
`error=reason`, bumps their lease to `now`, and returns the count. This is the restart
|
||||
reconciliation hook; `serve` calls it on startup (task 08).
|
||||
8. `BatchNbytes()` sums `nbytes` over `batch_job`. Batch bytes are **not** included in `Size()` or
|
||||
the content `evict()` union; batch is a coherent unit with its own budget.
|
||||
9. Accessors follow the existing `dcache` style: return `(value, bool)` or `error`, never panic,
|
||||
count a miss in the appropriate `*M` counter. Reuse the existing `ClassFleet`/`ClassRestricted`
|
||||
constants for `caller_class`.
|
||||
|
||||
## Given
|
||||
|
||||
Copy `docs/plans/p6/_files/dcache/batch_test.go` → `internal/dcache/batch_test.go`. It opens an in-memory `DB`
|
||||
and asserts the create/read round trip, `BatchJobByIDem` lease expiry, caller scoping and the
|
||||
status filter (rule 3), live-job
|
||||
counts, ordered targets, the exactly-once `targets_done`/lease-bump rules, `SweepBatch`,
|
||||
`EvictBatchOldest` (oldest-first, never a live job), `MarkInFlightFailed`, and `BatchNbytes`.
|
||||
|
||||
## Steps
|
||||
|
||||
- [ ] **1.** Branch; copy the given test. **2. See it fail** (`go test ./internal/dcache/` does not compile).
|
||||
- [ ] **3.** Schema in `Open`, the types, the methods, in a new `internal/dcache/batch.go`.
|
||||
- [ ] **4.** `gofmt -w`; `go test -race -count=3 ./internal/dcache/` → `ok` (every existing test too).
|
||||
- [ ] **5.** `make gate`. **6.** Row `p6/01-batchstore` in `docs/implementer-log.md`; commit.
|
||||
|
||||
```sh
|
||||
git add internal/dcache docs/implementer-log.md
|
||||
git commit -F .state/commit-msg.txt
|
||||
```
|
||||
|
||||
## Done when
|
||||
|
||||
- The given and existing dcache tests pass under `-race -count=3`; `make gate` ok; the given file is byte-identical.
|
||||
@@ -0,0 +1,199 @@
|
||||
# P6 task 02: the batch runner's core in `internal/batch` (no workers yet)
|
||||
|
||||
**Branch:** `p6` (`git switch p6`; `git status --short` must be empty, otherwise stop)
|
||||
**Commit subject:** `batch: runner core — types, Submit with idempotency and caps, Status, Touch, Cancel`
|
||||
|
||||
## Goal
|
||||
|
||||
A new package `internal/batch`. This task writes the whole public surface, and the part of it
|
||||
that needs no goroutines:
|
||||
- `Submit` (idempotency, caps, storing the job);
|
||||
- `Status`, `Touch`;
|
||||
- `Cancel` for a job no worker holds.
|
||||
|
||||
**Task 03 adds the workers.** Until then every job stays `queued`: `New` starts nothing, and
|
||||
`Close` has nothing to stop.
|
||||
|
||||
The runner knows **nothing** about HTTP, the guard, or the real fetch pipeline. The per-target
|
||||
work is an injected `FetchFunc` (unused until task 03), and persistence is the `*dcache.DB` from
|
||||
task 01.
|
||||
|
||||
## Files
|
||||
|
||||
- Copy: `docs/plans/p6/_files/batch/runner_core_test.go` → `internal/batch/runner_core_test.go`
|
||||
- Create: `internal/batch/batch.go`: every type, variable and function in Interfaces below.
|
||||
Task 03 adds `worker.go`.
|
||||
- Modify: `docs/implementer-log.md`
|
||||
|
||||
**Work in small steps, in files.** First write `batch.go` with the types and function
|
||||
signatures, and make it compile (`go vet ./internal/batch/`). Then fill in one function at a
|
||||
time, running the tests as you go. Do not write the implementation out in your reply before
|
||||
putting it in a file.
|
||||
|
||||
## What already exists (do not re-read these packages; this is their API)
|
||||
|
||||
```go
|
||||
// internal/dcache (task 01). Every time argument is unix seconds.
|
||||
type BatchJob struct{ ID, Caller, CallerClass, IDemKey, Status, Error string; TargetsTotal, TargetsDone int; CreatedAt, Lease, NBytes int64 }
|
||||
type BatchTarget struct{ JobID string; Idx int; URL, Options, Status string; Result, Error, Hint *string }
|
||||
func (d *DB) CreateBatchJob(j BatchJob) error
|
||||
func (d *DB) BatchJob(id string) (BatchJob, bool) // a pure read
|
||||
func (d *DB) BatchJobByIDem(caller, idemKey string, now int64, ttl time.Duration) (string, bool) // queued/running/done/partial only
|
||||
func (d *DB) LiveBatchJobs(caller string) int // queued + running
|
||||
func (d *DB) LiveBatchJobsTotal() int
|
||||
func (d *DB) PutBatchTargets(jobID string, targets []BatchTarget) error
|
||||
func (d *DB) BatchTargets(jobID string) ([]BatchTarget, bool) // in idx order
|
||||
func (d *DB) UpdateBatchTarget(jobID string, idx int, now int64, status string, result, errMsg, hint *string) error
|
||||
func (d *DB) SetBatchJobStatus(id, status, errMsg string, now int64) error
|
||||
func (d *DB) TouchBatchJob(id string, now int64) error
|
||||
func (d *DB) SweepBatch(now int64, ttl time.Duration) int
|
||||
func (d *DB) EvictBatchOldest(maxBytes int64) int
|
||||
const ClassFleet, ClassRestricted = "fleet", "restricted"
|
||||
|
||||
// internal/auth
|
||||
type Caller struct{ Name, Workspace, KeySHA256 string }
|
||||
```
|
||||
|
||||
## Interfaces (exact: the given test compiles against them)
|
||||
|
||||
```go
|
||||
package batch
|
||||
|
||||
type Options struct {
|
||||
Mode string `json:"mode,omitempty"`
|
||||
Question string `json:"question,omitempty"`
|
||||
Pages string `json:"pages,omitempty"`
|
||||
Grep string `json:"grep,omitempty"`
|
||||
JQ string `json:"jq,omitempty"`
|
||||
Context *bool `json:"context,omitempty"` // nil = /v1/fetch's default (on when there is a question)
|
||||
Recover bool `json:"recover,omitempty"`
|
||||
}
|
||||
|
||||
type Target struct {
|
||||
URL string `json:"url"`
|
||||
Options Options `json:"options"`
|
||||
}
|
||||
|
||||
// Result is a finished target: the /v1/fetch envelope's content, and every other envelope field
|
||||
// in Meta. It is stored JSON-encoded in the target's result column.
|
||||
type Result struct {
|
||||
Content string `json:"content"`
|
||||
Meta map[string]any `json:"meta,omitempty"`
|
||||
}
|
||||
|
||||
// TargetError is a failure with an agent-facing hint. Status is the HTTP status the same
|
||||
// failure gets from /v1/fetch; the runner ignores it (the metrics hook reads it). Any other
|
||||
// error is stored with its Error() text and no hint.
|
||||
type TargetError struct {
|
||||
Status int
|
||||
Msg, Hint string
|
||||
}
|
||||
|
||||
func (e *TargetError) Error() string { return e.Msg }
|
||||
|
||||
// FetchFunc does one target for the caller that submitted the job.
|
||||
type FetchFunc func(ctx context.Context, c *auth.Caller, t Target) (Result, error)
|
||||
|
||||
type Job struct {
|
||||
ID string
|
||||
Caller string
|
||||
Status string // queued | running | done | partial | failed | cancelled
|
||||
Error string
|
||||
TargetsTotal int
|
||||
TargetsDone int
|
||||
CreatedAt int64 // unix seconds
|
||||
}
|
||||
|
||||
type TargetStatus struct {
|
||||
Idx int
|
||||
URL string
|
||||
Status string // pending | running | done | failed
|
||||
Result *Result
|
||||
Error string
|
||||
Hint string
|
||||
}
|
||||
|
||||
// Hooks observe the runner for metrics. Any nil field is skipped.
|
||||
type Hooks struct {
|
||||
Job func(outcome string) // created, reused, done, partial, failed, cancelled
|
||||
Target func(err error, d time.Duration, resultBytes int) // once per stored target; err nil on success
|
||||
JobSeconds func(d time.Duration) // submit → terminal, once per job
|
||||
Depth func(live int) // live jobs, after every change
|
||||
Inflight func(n int) // targets being fetched, after every change
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
TTL time.Duration // 0 → 24h
|
||||
TargetsPerJob int // 0 → 32
|
||||
JobsPerCaller int // 0 → 1
|
||||
JobsTotal int // 0 → 8
|
||||
Workers int // 0 → 1; values above 2 are clamped to 2
|
||||
MaxBytes int64 // 0 → 64 MiB
|
||||
RetryWait time.Duration // 0 → 2s: the pause before a busy target is tried again
|
||||
RetryFor time.Duration // 0 → 10m: how long one target may keep coming back busy
|
||||
Hooks Hooks
|
||||
}
|
||||
|
||||
var (
|
||||
ErrNoTargets = errors.New("no targets")
|
||||
ErrTooManyTargets = errors.New("too many targets")
|
||||
ErrCallerBusy = errors.New("busy")
|
||||
ErrBusy = errors.New("jobs busy")
|
||||
ErrRetryLater = errors.New("retry later") // a FetchFunc wraps this when the target was only busy
|
||||
)
|
||||
|
||||
func New(db *dcache.DB, fetch FetchFunc, cfg Config, log *slog.Logger) *Runner // starts the workers
|
||||
func (r *Runner) Submit(ctx context.Context, c *auth.Caller, callerClass string, targets []Target) (id string, reused bool, err error)
|
||||
func (r *Runner) Cancel(id string) bool
|
||||
func (r *Runner) Status(id string) (Job, []TargetStatus, bool) // a pure read
|
||||
func (r *Runner) Touch(id string) // bumps the lease to now
|
||||
func (r *Runner) Close() // stops the workers; waits for an in-flight fetch to return
|
||||
```
|
||||
|
||||
## Rules
|
||||
|
||||
1. **Submit, in this order:**
|
||||
- `len(targets) == 0` → `ErrNoTargets`.
|
||||
- **Idempotency.** `idem_key` = hex SHA-256 of `c.Name + "\x00" + json.Marshal(targets)`.
|
||||
If `BatchJobByIDem(c.Name, key, now, TTL)` finds a job, return its id with `reused = true`:
|
||||
no new job, no caps, `Hooks.Job("reused")`.
|
||||
- **Caps:** more than `TargetsPerJob` → `ErrTooManyTargets`; `LiveBatchJobs(c.Name) >= JobsPerCaller`
|
||||
→ `ErrCallerBusy`; `LiveBatchJobsTotal() >= JobsTotal` → `ErrBusy`.
|
||||
- Store the job **`queued`** (random 16-byte hex id, `CreatedAt` = `Lease` = now) and its targets
|
||||
(`pending`, `Options` = the JSON of `t.Options`), **before returning**. Keep `c` in memory
|
||||
for this job: jobs never resume after a restart, so the caller is never persisted.
|
||||
- `Hooks.Job("created")`; `SweepBatch(now, TTL)`. (Task 03 also wakes a worker here.)
|
||||
2. **`New`** stores its arguments and applies the `Config` defaults. It starts no goroutine
|
||||
(task 03 starts the workers there). **`Close`** is safe to call more than once.
|
||||
3. **Cancel(id)** returns false for an unknown id or a job already terminal (`done`, `partial`,
|
||||
`failed`, `cancelled`). Otherwise it sets the job `cancelled` at once (its pending targets
|
||||
stay `pending`) and calls `Hooks.Job("cancelled")`. Keep the rest of `Cancel` simple;
|
||||
task 03 adds the in-flight target case.
|
||||
4. **Status(id)** maps the store rows; `Result` is decoded from the result column. It writes
|
||||
nothing. **Touch(id)** is `TouchBatchJob(id, now)`; the API calls it after an owner's read.
|
||||
5. Library code never panics; a store error is logged (`Warn`, with the job id, never a URL) and
|
||||
the job goes on. Log lines carry the job id, counts and statuses, **never a URL or host**.
|
||||
|
||||
|
||||
## Facts
|
||||
|
||||
- `runner_core_test.go` holds the helpers that task 03's test also uses (`newRunner`,
|
||||
`waitStatus`, `urls`, and the callers `alice`, `bob`, `carol`).
|
||||
- Without workers, a submitted job stays `queued`, so it counts as live for the caps.
|
||||
|
||||
## Steps
|
||||
|
||||
- [ ] **1.** `git switch p6`; copy the given test. **2. See it fail** (does not compile).
|
||||
- [ ] **3.** `batch.go`: types and signatures first, compile; then `Submit`, `Status`, `Touch`,
|
||||
`Cancel`.
|
||||
- [ ] **4.** `gofmt -w`; `go test -race -count=3 ./internal/batch/` → `ok`. **5.** `make gate`.
|
||||
- [ ] **6.** Row `p6/02-runner-core`; commit.
|
||||
|
||||
```sh
|
||||
git add internal/batch docs/implementer-log.md
|
||||
git commit -F .state/commit-msg.txt
|
||||
```
|
||||
|
||||
## Done when
|
||||
|
||||
- The given test passes under `-race -count=3`; `make gate` ok; the given file is byte-identical.
|
||||
@@ -0,0 +1,167 @@
|
||||
# P6 task 03: the batch runner's workers
|
||||
|
||||
**Branch:** `p6` (`git switch p6`; `git status --short` must be empty, otherwise stop)
|
||||
**Commit subject:** `batch: workers — round-robin, serial per job, busy retry, cancel in flight, hooks`
|
||||
|
||||
## Goal
|
||||
|
||||
Task 02 wrote the runner's types, `Submit`, `Status`, `Touch` and `Cancel`; every job stays
|
||||
`queued`. This task adds the workers that drain jobs: one target at a time per job, round-robin
|
||||
across jobs, with a retry for a target that was only busy.
|
||||
|
||||
## Files
|
||||
|
||||
- Copy: `docs/plans/p6/_files/batch/runner_test.go` → `internal/batch/runner_test.go`
|
||||
- Create: `internal/batch/worker.go` (the workers, the rotation, one target, the terminal status)
|
||||
- Modify: `internal/batch/batch.go` (`New` starts the workers; `Close` stops them; `Submit` wakes
|
||||
one; `Cancel` handles a job with a target in flight), `docs/implementer-log.md`
|
||||
|
||||
**Work in small steps, in files.** Write `worker.go` with the worker loop and a stub for one
|
||||
target, and make it compile. Then fill it in, running the tests as you go. Do not write the
|
||||
implementation out in your reply before putting it in a file.
|
||||
|
||||
## What already exists (task 02, in `internal/batch/batch.go`)
|
||||
|
||||
All the types in the task 02 file: `Options`, `Target`, `Result`, `TargetError`, `FetchFunc`,
|
||||
`Job`, `TargetStatus`, `Hooks`, `Config`, the errors (including `ErrRetryLater`), `Runner`,
|
||||
`New`, `Submit`, `Status`, `Touch`, `Cancel`, `Close`. Read `batch.go` for how the runner keeps
|
||||
its jobs and each job's caller; the dcache API it uses is listed in the task 02 file's
|
||||
"What already exists".
|
||||
|
||||
## Design (follow it; do not redesign)
|
||||
|
||||
The concurrency is decided here so you can write it straight into `worker.go`. Put the design into
|
||||
code, not into your reply.
|
||||
|
||||
**State**, added to `Runner` (keep task 02's fields):
|
||||
|
||||
```go
|
||||
mu sync.Mutex // guards everything below
|
||||
cond *sync.Cond // on mu
|
||||
live []*liveJob // the rotation, in submit order
|
||||
cursor int // index in live where the next scan starts
|
||||
closed bool
|
||||
ctx context.Context // the FetchFunc context; Close cancels it
|
||||
stop context.CancelFunc
|
||||
wg sync.WaitGroup // the workers
|
||||
inflight int
|
||||
```
|
||||
|
||||
```go
|
||||
type liveJob struct {
|
||||
id string
|
||||
caller *auth.Caller
|
||||
targets []Target // as submitted
|
||||
created time.Time
|
||||
next int // index of the next target to fetch
|
||||
running bool // a target of this job is in flight
|
||||
started bool // set to "running" in the store already
|
||||
cancelled bool
|
||||
retryAt time.Time // zero, or when the busy target may be tried again
|
||||
busySince time.Time // zero, or when the current target first came back busy
|
||||
done, failed int
|
||||
}
|
||||
```
|
||||
|
||||
**Submit** (task 02's code, plus): after storing, lock `mu`, append a `liveJob`, `Hooks.Depth(len(live))`,
|
||||
`cond.Broadcast()`, unlock.
|
||||
|
||||
**A worker** (`New` starts `cfg.Workers` of them with `wg.Add`):
|
||||
|
||||
```
|
||||
loop:
|
||||
lock mu
|
||||
j := pick(now) // scan live from cursor, wrapping once: the first job with !running,
|
||||
// !cancelled, next < len(targets), and now >= retryAt; set cursor to
|
||||
// the index after it
|
||||
if j == nil:
|
||||
if closed: unlock; return
|
||||
if some live job has a future retryAt: time.AfterFunc(until the earliest, cond.Broadcast)
|
||||
cond.Wait(); unlock-free continue (Wait re-locks; go back to loop start with mu held)
|
||||
j.running = true; inflight++; Hooks.Inflight(inflight)
|
||||
if !j.started: SetBatchJobStatus(j.id, "running", "", now); j.started = true
|
||||
unlock
|
||||
res, err := call(ctx, j.caller, j.targets[j.next]) // recovers a panic into "internal error"
|
||||
lock mu
|
||||
j.running = false; inflight--; Hooks.Inflight(inflight)
|
||||
if errors.Is(err, ErrRetryLater):
|
||||
if j.busySince is zero: j.busySince = now
|
||||
if now − busySince <= RetryFor: j.retryAt = now + RetryWait; Broadcast; unlock; continue
|
||||
// else fall through and store it failed with err.Error()
|
||||
store target j.next (rule 2); j.next++; j.busySince, j.retryAt = zero; count done/failed
|
||||
Hooks.Target(...); EvictBatchOldest(MaxBytes)
|
||||
if j.cancelled or j.next == len(targets):
|
||||
if !j.cancelled: set the terminal status (rule 3)
|
||||
remove j from live (fix cursor), forget its caller, Hooks.Depth(len(live))
|
||||
Broadcast; unlock
|
||||
```
|
||||
|
||||
Store calls happen under `mu`. They are short (one SQLite connection), and holding `mu` keeps
|
||||
the store consistent with the in-memory state.
|
||||
|
||||
**Cancel:** lock `mu`. If the job is in `live`: set `j.cancelled`, `SetBatchJobStatus(cancelled)`,
|
||||
and `Hooks.Job("cancelled")`. If `!j.running`, remove it from `live` now; a running one is
|
||||
removed by its worker after the store (above), which must not overwrite `cancelled`. Broadcast,
|
||||
unlock. A job not in `live` keeps task 02's behaviour.
|
||||
|
||||
**Close:** lock; `closed = true`; `stop()`; Broadcast; unlock; `wg.Wait()`. Safe to call twice.
|
||||
|
||||
A target stays `pending` in the store while it is in flight. The store has no `running` target
|
||||
status write, and no test needs one.
|
||||
|
||||
## Rules
|
||||
|
||||
1. **Workers.** `Workers` goroutines share one rotation of live jobs in submit order. A worker
|
||||
takes the next job in the rotation that has a pending target **and no target in flight**. It
|
||||
sets the job to `running` on its first target, fetches that one target, stores it, and moves
|
||||
on to the next job in the rotation. So:
|
||||
- each job is serial: its target `i+1` never starts before `i` is stored;
|
||||
- with one worker and two live jobs, their targets alternate (`a1 b1 a2 b2 …`);
|
||||
- a job that a worker has not reached yet stays `queued`.
|
||||
2. **One target:**
|
||||
- Call the `FetchFunc` with the job's caller, under a context that `Close` cancels. Recover
|
||||
a panic into the error `internal error` (never echo the panic value).
|
||||
- If the error wraps `ErrRetryLater`, the target stays `pending`: the worker goes on with the
|
||||
rotation, and the target is tried again after `RetryWait`. A target that has been coming
|
||||
back busy for longer than `RetryFor` is stored `failed` with the error's text.
|
||||
- Otherwise store it: success → `done` with the JSON of the `Result`; error → `failed`, with
|
||||
`TargetError`'s `Msg`/`Hint` or `err.Error()` and no hint.
|
||||
- Then call `Hooks.Target(err, time, len(result JSON))`, then `EvictBatchOldest(MaxBytes)`.
|
||||
3. **Terminal status.** When a job's last target is stored: all `done` → `done`; all `failed` →
|
||||
`failed`; a mix → `partial`. `SetBatchJobStatus`, then `Hooks.Job(status)` and
|
||||
`Hooks.JobSeconds(now − created)`, exactly once per job.
|
||||
4. **Cancel(id)** returns false for an unknown id or a job already terminal. Otherwise it sets the job
|
||||
`cancelled` at once. The in-flight target, if any, finishes and is stored; no further target
|
||||
starts; later targets stay `pending`. Storing that last target must not overwrite `cancelled`.
|
||||
`Hooks.Job("cancelled")` once.
|
||||
5. **Hooks.Depth** gets the live-job count and **Hooks.Inflight** the in-flight count whenever
|
||||
either changes.
|
||||
|
||||
## Facts
|
||||
|
||||
- The given test uses `FetchFunc`s that hold on channels, fail, return `ErrRetryLater`, or panic.
|
||||
It calls `Close` before the store closes.
|
||||
- `TestRoundRobin` holds the first target until the second job is submitted, then expects
|
||||
`a1 b1 a2 b2 a3 b3` from one worker.
|
||||
- Task 02's tests (`runner_core_test.go`) must still pass: with workers, a job there may
|
||||
already be `running` or `done` where it was `queued`; those tests allow for that.
|
||||
|
||||
## Steps
|
||||
|
||||
- [ ] **1.** `git switch p6`; copy the given test. **2. See it fail** (the jobs never finish).
|
||||
- [ ] **3.** `worker.go`. **4.** The `batch.go` changes.
|
||||
- [ ] **5.** `gofmt -w`; `go test -race -count=3 ./internal/batch/` → `ok`. **6.** `make gate`.
|
||||
- [ ] **7.** Row `p6/03-runner-workers`; commit.
|
||||
|
||||
```sh
|
||||
git add internal/batch docs/implementer-log.md
|
||||
git commit -F .state/commit-msg.txt
|
||||
```
|
||||
|
||||
## Done when
|
||||
|
||||
- Both given tests pass under `-race -count=3`; `make gate` ok; the given files are byte-identical.
|
||||
|
||||
## Stop and report if
|
||||
|
||||
- A given test's timing can only be met by sleeping longer than `RetryWait` inside a worker.
|
||||
@@ -0,0 +1,83 @@
|
||||
# P6 task 04: `sift_batch_*` series in `internal/metrics`
|
||||
|
||||
**Branch:** `p6` (`git switch p6`; `git status --short` must be empty, otherwise stop)
|
||||
**Commit subject:** `metrics: sift_batch_* series (jobs, targets, depth, in-flight, durations, bytes)`
|
||||
|
||||
Add the batch series to `internal/metrics` so the caps (32 targets, 1 job/caller, 8 jobs total)
|
||||
and the byte budget can be re-sized from observed data. **This task touches only
|
||||
`internal/metrics`**: task 06 wires these methods to the runner's `batch.Hooks`.
|
||||
|
||||
## Series (names and label values exact — the given test asserts these)
|
||||
|
||||
| Metric | Type | Labels / meaning |
|
||||
|---|---|---|
|
||||
| `sift_batch_jobs_total` | counter | `outcome` ∈ `created`, `reused`, `done`, `partial`, `failed`, `cancelled`, `restart` |
|
||||
| `sift_batch_targets_total` | counter | `outcome` ∈ `ok`, `blocked`, `timeout`, `busy`, `upstream_4xx`, `upstream_5xx`, `unsupported`, `too_long`, `cancel` |
|
||||
| `sift_batch_queue_depth` | gauge | live jobs (`queued`+`running`) |
|
||||
| `sift_batch_targets_inflight` | gauge | targets currently being fetched |
|
||||
| `sift_batch_job_seconds` | histogram | wall time from submit to terminal status |
|
||||
| `sift_batch_target_seconds` | histogram | wall time per target |
|
||||
| `sift_batch_result_bytes` | histogram | stored bytes per target result |
|
||||
|
||||
Classifying an error into a target outcome is task 06's job (it has the `*fetch.Error`); this
|
||||
package only records the label it is given.
|
||||
|
||||
## Public surface (exact names — the given test compiles against them)
|
||||
|
||||
```go
|
||||
// fields added to metrics.Set (registered in New)
|
||||
BatchJobsTotal *prometheus.CounterVec // label "outcome"
|
||||
BatchTargetsTotal *prometheus.CounterVec // label "outcome"
|
||||
BatchQueueDepth prometheus.Gauge
|
||||
BatchTargetsInflight prometheus.Gauge
|
||||
BatchJobSeconds prometheus.Histogram
|
||||
BatchTargetSeconds prometheus.Histogram
|
||||
BatchResultBytes prometheus.Histogram
|
||||
|
||||
// methods (nil-guarded, like ObservePDF)
|
||||
func (s *Set) ObserveBatchJob(outcome string)
|
||||
func (s *Set) ObserveBatchTarget(outcome string)
|
||||
func (s *Set) SetBatchQueueDepth(n int)
|
||||
func (s *Set) SetBatchTargetsInflight(n int)
|
||||
func (s *Set) ObserveBatchJobSeconds(d time.Duration)
|
||||
func (s *Set) ObserveBatchTargetSeconds(d time.Duration)
|
||||
func (s *Set) ObserveBatchResultBytes(n int)
|
||||
```
|
||||
|
||||
Both `outcome` counters fold any value outside their known set to `other` — the same "a raw string
|
||||
never becomes a label" guard as `ObservePDF`. Job outcomes: `created, reused, done, partial,
|
||||
failed, cancelled, restart`. Target outcomes: `ok, blocked, timeout, busy, upstream_4xx,
|
||||
upstream_5xx, unsupported, too_long, cancel`.
|
||||
|
||||
## Rules
|
||||
|
||||
1. A new file `internal/metrics/batch.go` adds the series to `metrics.Set` (registered in `New`),
|
||||
with `Observe…`/`Set…` methods mirroring `ObservePDF`/`ObserveDigest`, nil-receiver safe.
|
||||
2. Both `outcome` counters fold any value outside their known set to `other`: a raw string
|
||||
never becomes a label.
|
||||
3. The series are exposed through the existing `/metrics` and `/metrics.json` paths: no new
|
||||
endpoint. `Snapshot`/`writeMetricsJSON` do not need to list them.
|
||||
4. Histogram buckets: seconds as the existing request histograms; bytes
|
||||
`1K, 4K, 16K, 64K, 256K, 1M`.
|
||||
|
||||
## Given
|
||||
|
||||
Copy `docs/plans/p6/_files/metrics/batch_test.go` → `internal/metrics/batch_test.go`. It drives the `Observe…`
|
||||
methods, serves the `/metrics` handler (as `pdf_test.go` does), and asserts the exact series names,
|
||||
label values, gauge values, histogram counts, and that a raw outcome string is folded to `other`.
|
||||
|
||||
## Steps
|
||||
|
||||
- [ ] **1.** `git switch p6`; copy the given test. **2. See it fail.**
|
||||
- [ ] **3.** `internal/metrics/batch.go` and the registration in `New`.
|
||||
- [ ] **4.** `gofmt -w`; `go test -race -count=3 ./internal/metrics/` → `ok`. **5.** `make gate`.
|
||||
- [ ] **6.** Row `p6/04-metrics`; commit.
|
||||
|
||||
```sh
|
||||
git add internal/metrics docs/implementer-log.md
|
||||
git commit -F .state/commit-msg.txt
|
||||
```
|
||||
|
||||
## Done when
|
||||
|
||||
- The given and existing metrics tests pass under `-race -count=3`; `make gate` ok; the given file is byte-identical.
|
||||
@@ -0,0 +1,86 @@
|
||||
# P6 task 05: extract `v1fetch`'s pipeline into `fetchTarget`
|
||||
|
||||
**Branch:** `p6` (`git switch p6`; `git status --short` must be empty, otherwise stop)
|
||||
**Commit subject:** `api: v1fetch's per-target pipeline becomes fetchTarget (no behaviour change)`
|
||||
|
||||
> This is the one task in P6 that edits working code. It is a **move, not a rewrite**: the
|
||||
> reviewer will diff the old `v1fetch` body against `fetchTarget` line by line.
|
||||
|
||||
## Goal
|
||||
|
||||
`v1fetch` today inlines the whole per-target pipeline: fetch → recover → PDF render or clean →
|
||||
question limit → digest → context → log line → envelope. The batch worker (task 06) needs the
|
||||
same pipeline for one target. Move it into one method that both call:
|
||||
|
||||
```go
|
||||
// internal/api/target.go
|
||||
func (s *Server) fetchTarget(ctx context.Context, c *auth.Caller, t batch.Target) (batch.Result, error)
|
||||
```
|
||||
|
||||
`v1fetch` keeps its exact responses. **Every existing api test must pass unchanged.**
|
||||
|
||||
## Files
|
||||
|
||||
- Copy: `docs/plans/p6/_files/api/fetchtarget_test.go` → `internal/api/fetchtarget_test.go`
|
||||
- Create: `internal/api/target.go` (`fetchTarget`)
|
||||
- Modify: `internal/api/api.go` (`v1fetch` shrinks; `doFetch` takes a `context.Context` in place of
|
||||
the `*http.Request`, and its call sites pass `r.Context()`), `docs/implementer-log.md`
|
||||
|
||||
## What already exists
|
||||
|
||||
```go
|
||||
// internal/batch (task 02)
|
||||
type Options struct{ Mode, Question, Pages, Grep, JQ string; Context *bool; Recover bool }
|
||||
type Target struct{ URL string; Options Options }
|
||||
type Result struct{ Content string; Meta map[string]any }
|
||||
```
|
||||
|
||||
`fetch.Error{Status int; Msg, Hint string; RetryAfter int}` is the error type `writeFetchErr`
|
||||
already turns into a response.
|
||||
|
||||
## Rules
|
||||
|
||||
1. **What stays in `v1fetch`:** `s.reqN.Add(1)`, the caller check, decoding `v1Req`, and the
|
||||
`400 invalid url` for a bad body or an empty URL. Then it builds a `batch.Target` from the
|
||||
request (`URL`, and `Mode`, `Question`, `Pages`, `Grep`, `JQ`, `Context`, `Recover` into
|
||||
`Options`), calls `s.fetchTarget(r.Context(), c, t)`, and:
|
||||
- on an error, `s.writeFetchErr(w, c, req.URL, err)`;
|
||||
- on success, `writeJSON(w, 200, env)` where `env` is `res.Meta` plus `"content": res.Content`.
|
||||
2. **What moves into `fetchTarget`, in the same order and unchanged:** the question truncation
|
||||
to 500 runes, the mode default and check, `wantCtx`, `doFetch`, the recover path
|
||||
(`recover.Try(ctx, …)`), the PDF branch (render, `pages`/`grep` errors, the question-chunk
|
||||
limit), the clean/raw branch, the digest, the context search, the `fetch` log line, and the
|
||||
envelope.
|
||||
3. **Errors that `v1fetch` wrote directly with `writeErr` become `*fetch.Error` values with the same
|
||||
status, message and hint:** `400 invalid url` for a bad mode, `400 raw mode does not apply to
|
||||
PDFs`. `writeFetchErr` then writes them exactly as before. The only visible difference is its
|
||||
`fetch failed` log line, which is fine.
|
||||
4. **The envelope:** `fetchTarget` returns `Result{Content: content, Meta: m}` where `m` holds every
|
||||
other key of today's `out` map (`mode`, `untrusted`, `source`, `warnings`, `sizes`, `timings`,
|
||||
and `pdf`, `digest`, `context` when present). Never put `content` in `Meta`.
|
||||
5. **`doFetch(ctx, c, rawURL, jq, onlyMain)`:** `r.Context()` becomes `ctx`. That is the only change
|
||||
in `doFetch`. `scrape` (the other caller) passes `r.Context()`.
|
||||
6. No new behaviour, no new log fields, no renamed JSON keys. If something cannot move without
|
||||
changing behaviour, stop and report it.
|
||||
|
||||
## Steps
|
||||
|
||||
- [ ] **1.** `git switch p6`; copy the given test. **2. See it fail** (`fetchTarget` undefined).
|
||||
- [ ] **3.** `doFetch(ctx, …)` and its call sites. **4.** `target.go`. **5.** Shrink `v1fetch`.
|
||||
- [ ] **6.** `gofmt -w`; `go test -race -count=3 ./internal/api/` → `ok`, every existing test
|
||||
included. **7.** `make gate`.
|
||||
- [ ] **8.** Row `p6/05-fetchtarget`; commit.
|
||||
|
||||
```sh
|
||||
git add internal/api docs/implementer-log.md
|
||||
git commit -F .state/commit-msg.txt
|
||||
```
|
||||
|
||||
## Done when
|
||||
|
||||
- The two given tests and **every existing api test** pass under `-race -count=3`, with no
|
||||
existing test file modified; `make gate` ok; the given file is byte-identical.
|
||||
|
||||
## Stop and report if
|
||||
|
||||
- An existing api test fails and fixing it would need a change to that test.
|
||||
@@ -0,0 +1,129 @@
|
||||
# P6 task 06: the runner wired to `fetchTarget` and the metrics
|
||||
|
||||
**Branch:** `p6` (`git switch p6`; `git status --short` must be empty, otherwise stop)
|
||||
**Commit subject:** `api: NewBatchRunner — the batch fetch adapter, retry on busy, outcome metrics`
|
||||
|
||||
## Goal
|
||||
|
||||
Build the production runner: its `FetchFunc` is an adapter over `fetchTarget` (task 05), and its
|
||||
hooks feed the `sift_batch_*` metrics (task 04). **No HTTP endpoints yet** (task 07).
|
||||
|
||||
## Files
|
||||
|
||||
- Copy: `docs/plans/p6/_files/api/batchrun_test.go` → `internal/api/batchrun_test.go`
|
||||
- Create: `internal/api/batch.go` (`NewBatchRunner`, `batchFetch`, `batchOutcome`)
|
||||
- Modify:
|
||||
- `internal/api/api.go`: only the `Batch` field.
|
||||
- `internal/api/pdf.go`: only the pool caller name (rule 3).
|
||||
- `docs/implementer-log.md`
|
||||
|
||||
## What already exists
|
||||
|
||||
```go
|
||||
// internal/batch (task 02)
|
||||
type Options struct{ Mode, Question, Pages, Grep, JQ string; Context *bool; Recover bool }
|
||||
type Target struct{ URL string; Options Options }
|
||||
type Result struct{ Content string; Meta map[string]any }
|
||||
type TargetError struct{ Status int; Msg, Hint string }
|
||||
type FetchFunc func(ctx context.Context, c *auth.Caller, t Target) (Result, error)
|
||||
type Job struct{ ID, Caller, Status, Error string; TargetsTotal, TargetsDone int; CreatedAt int64 }
|
||||
type TargetStatus struct{ Idx int; URL, Status string; Result *Result; Error, Hint string }
|
||||
type Hooks struct {
|
||||
Job func(outcome string)
|
||||
Target func(err error, d time.Duration, resultBytes int)
|
||||
JobSeconds func(d time.Duration)
|
||||
Depth func(live int)
|
||||
Inflight func(n int)
|
||||
}
|
||||
type Config struct{ TTL time.Duration; TargetsPerJob, JobsPerCaller, JobsTotal, Workers int; MaxBytes int64; RetryWait, RetryFor time.Duration; Hooks Hooks }
|
||||
var ErrNoTargets, ErrTooManyTargets, ErrCallerBusy, ErrBusy, ErrRetryLater error
|
||||
func New(db *dcache.DB, fetch FetchFunc, cfg Config, log *slog.Logger) *Runner
|
||||
func (r *Runner) Submit(ctx context.Context, c *auth.Caller, callerClass string, targets []Target) (string, bool, error)
|
||||
func (r *Runner) Cancel(id string) bool
|
||||
func (r *Runner) Status(id string) (Job, []TargetStatus, bool)
|
||||
func (r *Runner) Touch(id string)
|
||||
func (r *Runner) Close()
|
||||
|
||||
// internal/metrics (task 04), on s.Prom
|
||||
ObserveBatchJob(outcome string); ObserveBatchTarget(outcome string)
|
||||
SetBatchQueueDepth(n int); SetBatchTargetsInflight(n int)
|
||||
ObserveBatchJobSeconds(d time.Duration); ObserveBatchTargetSeconds(d time.Duration); ObserveBatchResultBytes(n int)
|
||||
|
||||
// internal/api (task 05 and before)
|
||||
func (s *Server) fetchTarget(ctx context.Context, c *auth.Caller, t batch.Target) (batch.Result, error)
|
||||
func (s *Server) caller(w http.ResponseWriter, r *http.Request) *auth.Caller // nil after writing 401/403
|
||||
func (s *Server) callerClass(c *auth.Caller) string
|
||||
func (s *Server) logErr(c *auth.Caller, err error) any
|
||||
// requestDeadline (210 s), maxReqBody (64 KiB), writeErr, writeJSON
|
||||
```
|
||||
|
||||
## Interfaces added here
|
||||
|
||||
```go
|
||||
// in Server (after pdfPoolOnce):
|
||||
Batch *batch.Runner // nil: the /v1/batch endpoints (task 07) answer 503 "batch disabled"
|
||||
|
||||
// NewBatchRunner builds the runner that serve (task 08) and the tests use: its FetchFunc is
|
||||
// batchFetch and its Hooks feed s.Prom. cfg.Hooks is overwritten.
|
||||
func (s *Server) NewBatchRunner(db *dcache.DB, cfg batch.Config) *batch.Runner
|
||||
|
||||
func (s *Server) batchFetch(ctx context.Context, c *auth.Caller, t batch.Target) (batch.Result, error)
|
||||
func batchOutcome(err error) string
|
||||
```
|
||||
|
||||
## Rules
|
||||
|
||||
1. **`batchFetch`**, the runner's `FetchFunc`:
|
||||
- `context.WithTimeout(ctx, requestDeadline)`, mark the context as a batch fetch (an unexported
|
||||
context key), and call `s.fetchTarget`.
|
||||
- Map its error:
|
||||
- a `*fetch.Error` with status `429` or `503` → `fmt.Errorf("%s: %w", fe.Msg, batch.ErrRetryLater)`,
|
||||
so the runner retries it;
|
||||
- any other `*fetch.Error` → `&batch.TargetError{Status, Msg, Hint}`, with `502` rewritten to
|
||||
`500 internal error` as `writeFetchErr` does, and `s.blockN.Add(1)` for a `403`;
|
||||
- anything else → `&batch.TargetError{Status: 500, Msg: "internal error"}`, logged `Error`
|
||||
with `s.logErr(c, err)`.
|
||||
2. **Hooks** (in `NewBatchRunner`):
|
||||
- `Job` → `ObserveBatchJob`;
|
||||
- `Target(err, d, n)` → `ObserveBatchTarget(batchOutcome(err))`, `ObserveBatchTargetSeconds(d)`, and
|
||||
`ObserveBatchResultBytes(n)` when `err == nil`;
|
||||
- `JobSeconds`, `Depth`, `Inflight` → the matching `s.Prom` methods.
|
||||
- `batchOutcome(err)` is:
|
||||
- `ok` for nil;
|
||||
- `busy` for an error wrapping `ErrRetryLater` (the runner gave up retrying);
|
||||
- otherwise, by `TargetError.Status`: `403` `blocked`, `429` `busy`, `504` `timeout`,
|
||||
`415` `unsupported`, `422` whose message starts `document too long` `too_long`, other
|
||||
4xx `upstream_4xx`, 5xx `upstream_5xx`; any other error `upstream_5xx`.
|
||||
3. **The PDF pool under a batch:** in `extractPDF`, the pool's caller argument is `c.Name+"/batch"`
|
||||
when the context carries the batch mark, else `c.Name` as today. A batch extraction then
|
||||
never takes the caller's own interactive slot. The single-flight key is unchanged, so a batch
|
||||
and an interactive extraction of the same PDF still share one child.
|
||||
4. Log lines added here carry the caller and the job id, **never a target URL or hostname**. The
|
||||
per-target `fetch` line inside `fetchTarget` is unchanged.
|
||||
|
||||
## Facts
|
||||
|
||||
- The given api test builds the server through `s.NewBatchRunner` over an in-memory `dcache`
|
||||
(`newBatchServer`). Targets are local `httptest` servers, or `internal.test`, which the test's
|
||||
`LookupIP` resolves to `10.0.0.5`.
|
||||
- `TestBatchFetchErrors` fills the 8-slot global gate with `s.acquire` so the next fetch is
|
||||
refused busy after `acquire`'s 5 s wait.
|
||||
- `batchrun_test.go` also holds the helpers task 07's test uses; it is a given file like the others.
|
||||
|
||||
## Steps
|
||||
|
||||
- [ ] **1.** `git switch p6`; copy the given test. **2. See it fail.**
|
||||
- [ ] **3.** `batch.go`. **4.** The `Batch` field and the pool name.
|
||||
- [ ] **5.** `gofmt -w`; `go test -race -count=3 ./internal/api/` → `ok`.
|
||||
**6.** `make gate`.
|
||||
- [ ] **7.** Row `p6/06-batch-wiring`; commit.
|
||||
|
||||
```sh
|
||||
git add internal/api docs/implementer-log.md
|
||||
git commit -F .state/commit-msg.txt
|
||||
```
|
||||
|
||||
## Done when
|
||||
|
||||
- The given test and every existing api test pass under `-race -count=3`; `make gate` ok; the
|
||||
given file is byte-identical.
|
||||
@@ -0,0 +1,112 @@
|
||||
# P6 task 07: the `/v1/batch` endpoints
|
||||
|
||||
**Branch:** `p6` (`git switch p6`; `git status --short` must be empty, otherwise stop)
|
||||
**Commit subject:** `api: /v1/batch — submit, status summary, one result, cancel`
|
||||
|
||||
## Goal
|
||||
|
||||
The HTTP surface over `s.Batch` (task 06): submit, a status summary, one target's result, and
|
||||
cancel, with the submit-time guard pre-check, the caps as HTTP statuses, and ownership.
|
||||
|
||||
## Files
|
||||
|
||||
- Copy: `docs/plans/p6/_files/api/batch_test.go` → `internal/api/batch_test.go`
|
||||
- Modify: `internal/api/batch.go` (the four handlers), `internal/api/api.go` (only the routes in
|
||||
`New` and `ServeHTTP`, rule 2), `docs/implementer-log.md`
|
||||
|
||||
## What already exists
|
||||
|
||||
```go
|
||||
// internal/batch (task 02)
|
||||
type Target struct{ URL string; Options Options }
|
||||
type Options struct{ Mode, Question, Pages, Grep, JQ string; Context *bool; Recover bool }
|
||||
type Result struct{ Content string; Meta map[string]any }
|
||||
type Job struct{ ID, Caller, Status, Error string; TargetsTotal, TargetsDone int; CreatedAt int64 }
|
||||
type TargetStatus struct{ Idx int; URL, Status string; Result *Result; Error, Hint string }
|
||||
var ErrNoTargets, ErrTooManyTargets, ErrCallerBusy, ErrBusy error
|
||||
func (r *Runner) Submit(ctx context.Context, c *auth.Caller, callerClass string, targets []Target) (string, bool, error)
|
||||
func (r *Runner) Cancel(id string) bool
|
||||
func (r *Runner) Status(id string) (Job, []TargetStatus, bool) // a pure read
|
||||
func (r *Runner) Touch(id string) // bumps the lease to now
|
||||
|
||||
// internal/guard (on master)
|
||||
func (g *Guard) PreCheck(raw string) error // "blocked: …" errors
|
||||
|
||||
// internal/api
|
||||
Batch *batch.Runner // in Server (task 06); nil → 503 "batch disabled"
|
||||
func (s *Server) caller(w http.ResponseWriter, r *http.Request) *auth.Caller // nil after writing 401/403
|
||||
func (s *Server) callerClass(c *auth.Caller) string
|
||||
// maxReqBody (64 KiB), writeErr(w, status, msg, hint), writeJSON(w, status, v), s.Prom.Observe(path, d, code)
|
||||
```
|
||||
|
||||
Add `const batchMaxBody = 256 << 10`.
|
||||
|
||||
## Rules
|
||||
|
||||
1. **Routes** (in `New`): `POST /v1/batch`, `GET /v1/batch/{id}`, `GET /v1/batch/{id}/{idx}`,
|
||||
`DELETE /v1/batch/{id}`. Each starts with `s.caller(w, r)`.
|
||||
2. **`ServeHTTP`:**
|
||||
- The body cap is `batchMaxBody` for `/v1/batch` and `maxReqBody` for every other path. The cap
|
||||
is applied once, here: a handler cannot raise a cap already set.
|
||||
- Batch paths are observed with a fixed label: `/v1/batch` or `/v1/batch/{id}`. A job id must
|
||||
never become a label value.
|
||||
3. **Submit** (`POST /v1/batch`), checked in this order:
|
||||
- The body is `{"targets": [{"url", "mode", "question", "pages", "grep", "jq", "context",
|
||||
"recover"}, …]}`, each target flat with `/v1/fetch`'s field names. Bad JSON, a missing or
|
||||
empty list, an empty `url`, or a `mode` other than empty, `clean`, `raw` or `digest` →
|
||||
`400 invalid batch: …`.
|
||||
- `s.Guard.PreCheck(url)` for every target. The first failure → `422` with the message
|
||||
`blocked: target <idx>: <reason>`, and no job is created.
|
||||
- `s.Batch.Submit(r.Context(), c, s.callerClass(c), targets)`:
|
||||
- `ErrTooManyTargets` → `422 too many targets (max 32)` (the configured cap);
|
||||
- `ErrCallerBusy` → `429 busy: a batch of yours is still running`, `Retry-After: 30`;
|
||||
- `ErrBusy` → `429 batch jobs busy`, `Retry-After: 60`;
|
||||
- `ErrNoTargets` → `400`.
|
||||
- `202 {"success": true, "job_id", "status", "targets": n, "reused"}`. `status` is `"queued"`
|
||||
for a new job, and the job's current status for a reused one.
|
||||
4. **Ownership:** a job whose `Caller` is not `c.Name`, or an unknown id, is `404 not found` for
|
||||
every GET and DELETE. It is never `403`: a caller cannot probe for another caller's jobs.
|
||||
5. **Status** (`GET /v1/batch/{id}`): after the ownership check, `s.Batch.Touch(id)`, then
|
||||
`200 {"job_id", "status", "error" (only when set), "targets_total", "targets_done",
|
||||
"created_at" (RFC 3339, UTC), "targets": [{"idx", "url", "status", "error"?, "hint"?}]}`.
|
||||
- This is a summary: **no results**, because 32 inline results can be 6 MB, too much for one
|
||||
agent tool call.
|
||||
- With `?results=1`, each `done` target also carries `"result"`: the envelope, i.e. `Meta`
|
||||
plus `"content"`.
|
||||
6. **One result** (`GET /v1/batch/{id}/{idx}`): after ownership and `Touch`, `200 {"job_id", "idx",
|
||||
"url", "status", "result"?, "error"?, "hint"?}`. An `idx` that is not a number in
|
||||
`0..targets_total-1` → `404`. **Cancel** (`DELETE`): after ownership, a terminal job →
|
||||
`409 job already finished`; otherwise `Cancel` → `200 {"success": true, "status": "cancelled"}`.
|
||||
7. **Logging:** each batch handler logs one line: `batch` with `caller`, `job_id`, `targets` and
|
||||
`status` (or `reused`). **Never a target URL or hostname, for any caller**, and never the
|
||||
`idem_key`. The per-target `fetch` line inside `fetchTarget` is unchanged: it is redacted
|
||||
for restricted callers exactly as for `/v1/fetch`.
|
||||
|
||||
|
||||
## Facts
|
||||
|
||||
- `newBatchServer`, `pageServer` and `fakeDNS` come from `batchrun_test.go` (task 06). Targets
|
||||
are local `httptest` servers, some held on a channel, or `internal.test` (resolves to `10.0.0.5`).
|
||||
- `TestBatchBodyCap` posts 32 targets with questions, about 66 KiB: over `maxReqBody`, under
|
||||
`batchMaxBody`.
|
||||
|
||||
## Steps
|
||||
|
||||
- [ ] **1.** `git switch p6`; copy the given test. **2. See it fail.**
|
||||
- [ ] **3.** The handlers. **4.** Routes and `ServeHTTP`.
|
||||
- [ ] **5.** `gofmt -w`; `go test -race -count=3 ./internal/api/` → `ok`. **6.** `make gate`.
|
||||
- [ ] **7.** Row `p6/07-api-endpoints`; commit.
|
||||
|
||||
```sh
|
||||
git add internal/api docs/implementer-log.md
|
||||
git commit -F .state/commit-msg.txt
|
||||
```
|
||||
|
||||
## Done when
|
||||
|
||||
- The given test and every existing api test pass under `-race -count=3`; `make gate` ok; the
|
||||
given file is byte-identical.
|
||||
|
||||
## Stop and report if
|
||||
|
||||
- Any given test needs a URL that is not local, or a change to an existing test.
|
||||
@@ -0,0 +1,101 @@
|
||||
# P6 task 08: the `[batch]` config and `serve` wiring
|
||||
|
||||
**Branch:** `p6` (`git switch p6`; `git status --short` must be empty, otherwise stop)
|
||||
**Commit subject:** `serve: [batch] config; open the store, reap jobs killed by a restart, start the runner`
|
||||
|
||||
## Goal
|
||||
|
||||
`sift serve` runs the batch runner from `[batch]` config. On startup it marks jobs that a restart
|
||||
killed as failed (they are never resumed). On shutdown it stops the runner before closing the store.
|
||||
|
||||
## Files
|
||||
|
||||
- Copy: `docs/plans/p6/_files/config/batch_test.go` → `internal/config/batch_test.go`
|
||||
- Modify: `internal/config/config.go` (the `Batch` type and field, `ParseSize`),
|
||||
`cmd/sift/serve.go` (rule 3), `dist/sift.toml.example` (rule 4), `docs/implementer-log.md`
|
||||
|
||||
## What already exists
|
||||
|
||||
```go
|
||||
// internal/batch (task 02)
|
||||
type Config struct{ TTL time.Duration; TargetsPerJob, JobsPerCaller, JobsTotal, Workers int; MaxBytes int64; RetryWait, RetryFor time.Duration; Hooks Hooks }
|
||||
func (r *Runner) Close()
|
||||
|
||||
// internal/api (task 06)
|
||||
func (s *Server) NewBatchRunner(db *dcache.DB, cfg batch.Config) *batch.Runner
|
||||
Batch *batch.Runner // in Server
|
||||
// s.Prom.ObserveBatchJob(outcome string) (task 04)
|
||||
|
||||
// internal/dcache (task 01)
|
||||
func (d *DB) MarkInFlightFailed(reason string, now int64) int
|
||||
|
||||
// internal/config
|
||||
func ParseDuration(s string, def time.Duration) time.Duration
|
||||
```
|
||||
|
||||
## Interfaces added here
|
||||
|
||||
```go
|
||||
package config
|
||||
|
||||
type Batch struct {
|
||||
Disabled bool `toml:"disabled"`
|
||||
TargetsPerJob int `toml:"targets_per_job"`
|
||||
JobsPerCaller int `toml:"jobs_per_caller"`
|
||||
JobsTotal int `toml:"jobs_total"`
|
||||
Workers int `toml:"workers"`
|
||||
TTL string `toml:"ttl"`
|
||||
MaxBytes string `toml:"max_bytes"`
|
||||
}
|
||||
|
||||
// in File, after Search:
|
||||
Batch Batch `toml:"batch"`
|
||||
|
||||
// ParseSize parses a positive byte count: a plain integer, or one with KiB, MiB or GiB.
|
||||
// Anything else, zero or negative returns def.
|
||||
func ParseSize(s string, def int64) int64
|
||||
```
|
||||
|
||||
## Rules
|
||||
|
||||
1. Batch is **on by default**: no `[batch]` block means the runner with its own defaults.
|
||||
`disabled = true` turns it off, and the endpoints then answer 503.
|
||||
2. `ParseSize` trims spaces and accepts only the binary suffixes `KiB`, `MiB`, `GiB`, or none (bytes).
|
||||
3. **`serve.go`:**
|
||||
- Open the `dcache` store when batch is on, as well as in today's cases: same path, same call.
|
||||
If the store cannot be opened, log `Warn` and run without batch.
|
||||
- Then, **before** building the runner: `n := cache.MarkInFlightFailed("restart", now)`, and call
|
||||
`srv.Prom.ObserveBatchJob("restart")` `n` times. Log `batch reaped` with `n` if `n > 0`.
|
||||
- `srv.Batch = srv.NewBatchRunner(cache, batch.Config{...})`, with the config fields: `TTL` via
|
||||
`ParseDuration(…, 0)` and `MaxBytes` via `ParseSize(…, 0)`. The runner's defaults fill the zeros.
|
||||
- On shutdown (both branches of the final `select`), `srv.Batch.Close()` before the process
|
||||
exits, when it is non-nil.
|
||||
4. **`dist/sift.toml.example`:** a commented `[batch]` block with the defaults, in the style of the
|
||||
file:
|
||||
```toml
|
||||
# [batch]
|
||||
# disabled = false
|
||||
# targets_per_job = 32
|
||||
# jobs_per_caller = 1
|
||||
# jobs_total = 8
|
||||
# workers = 1 # at most 2: each worker fetches one target at a time
|
||||
# ttl = "24h" # a job's results live this long after its last read or write
|
||||
# max_bytes = "64MiB" # stored results; the oldest finished jobs go first
|
||||
```
|
||||
|
||||
## Steps
|
||||
|
||||
- [ ] **1.** `git switch p6`; copy the given test. **2. See it fail.**
|
||||
- [ ] **3.** `config.go`. **4.** `serve.go`. **5.** The example config.
|
||||
- [ ] **6.** `gofmt -w`; `go test -race -count=3 ./internal/config/ ./cmd/sift/` → `ok`. **7.** `make gate`.
|
||||
- [ ] **8.** Row `p6/08-serve-config`; commit.
|
||||
|
||||
```sh
|
||||
git add internal/config cmd/sift/serve.go dist/sift.toml.example docs/implementer-log.md
|
||||
git commit -F .state/commit-msg.txt
|
||||
```
|
||||
|
||||
## Done when
|
||||
|
||||
- The given test and every existing config and cmd test pass under `-race -count=3`; `make gate`
|
||||
ok; the given file is byte-identical.
|
||||
@@ -0,0 +1,89 @@
|
||||
# P6 task 09: `sift batch`, README and skill
|
||||
|
||||
**Branch:** `p6` (`git switch p6`; `git status --short` must be empty, otherwise stop)
|
||||
**Commit subject:** `cli: sift batch submit|status|get|cancel; README and skill`
|
||||
|
||||
## Goal
|
||||
|
||||
Agents reach the batch endpoint through the CLI, as they do `fetch` and `search`. There is no
|
||||
Hermes plugin.
|
||||
|
||||
## Files
|
||||
|
||||
- Copy: `docs/plans/p6/_files/cmd_sift/batch_test.go` → `cmd/sift/batch_test.go`
|
||||
- Create: `cmd/sift/batchcmd.go` (`batchMain`, `parseBatchArgs`, one small function per verb)
|
||||
- Modify: `cmd/sift/main.go` (one `case "batch":` line and the usage text), `README.md`,
|
||||
`skill/SKILL.md`, `docs/implementer-log.md`
|
||||
|
||||
## What already exists
|
||||
|
||||
- `cmd/sift`: `loadClient() (base, key string, err error)` reads `client.toml`; `searchcmd.go` is the
|
||||
pattern for a verb (build the request with the bearer header, decode the JSON envelope, exit
|
||||
non-zero on a non-2xx). Never put a key on the command line.
|
||||
- The endpoints (task 07):
|
||||
- `POST /v1/batch` `{"targets":[…]}` → `202 {"success","job_id","status","targets","reused"}`;
|
||||
- `GET /v1/batch/{id}` (summary; `?results=1` inlines results);
|
||||
- `GET /v1/batch/{id}/{idx}` (one result);
|
||||
- `DELETE /v1/batch/{id}`.
|
||||
|
||||
## Interfaces added here
|
||||
|
||||
```go
|
||||
// parseBatchArgs validates a `sift batch` command line (without "batch").
|
||||
func parseBatchArgs(args []string) (verb string, rest []string, err error)
|
||||
```
|
||||
|
||||
```
|
||||
sift batch submit <targets.json|-> # prints {"job_id", "status", "reused"}
|
||||
sift batch status [--results] <job_id> # prints the status body (a summary unless --results)
|
||||
sift batch get <job_id> <idx> # prints one target's result, as `sift fetch --json` would
|
||||
sift batch cancel <job_id> # prints {"success", "status"}
|
||||
```
|
||||
|
||||
## Rules
|
||||
|
||||
1. `parseBatchArgs` accepts exactly the forms above:
|
||||
- `submit` with one argument (`-` means stdin);
|
||||
- `status` with a job id, optionally preceded by `--results`;
|
||||
- `get` with a job id and an `idx` that is a non-negative integer;
|
||||
- `cancel` with a job id.
|
||||
Anything else (no verb, a missing or extra argument, an unknown verb, a bad `idx`) is an
|
||||
error. `rest` is the arguments after the verb, unchanged.
|
||||
2. `submit` reads the file and checks it locally before calling the server: it must be a JSON
|
||||
object with a non-empty `targets` array. It may also be a bare array of targets or of URL
|
||||
strings; wrap either into `{"targets": […]}`, turning each string into `{"url": s}`.
|
||||
3. Every verb prints the server's JSON on stdout, and on a non-2xx prints the error envelope to
|
||||
stderr and exits with `exitFor(status, msg)`, as `fetch` does.
|
||||
4. **`README.md`**, a short "Batch" section:
|
||||
- the four endpoints;
|
||||
- the model: serial per job, one worker by default, round-robin across jobs;
|
||||
- the caps (32 targets, 1 live job per caller, 8 in all) and the 256 KiB body;
|
||||
- idempotency: the same list from the same caller returns the same job while it is live or
|
||||
has results;
|
||||
- the lease/TTL and restart semantics;
|
||||
- the summary-vs-results reads;
|
||||
- restricted-caller redaction;
|
||||
- the `sift_batch_*` metrics.
|
||||
5. **`skill/SKILL.md`**, a short section:
|
||||
- for several URLs you will read later, `sift batch submit` them once, rather than a loop of
|
||||
`sift fetch`;
|
||||
- poll with `sift batch status`, and read each result with `sift batch get <job> <idx>`;
|
||||
- don't use `--results` unless you need every result at once (it can be large);
|
||||
- a batch is slow on purpose: for one URL you need now, use `sift fetch`.
|
||||
|
||||
## Steps
|
||||
|
||||
- [ ] **1.** `git switch p6`; copy the given test. **2. See it fail.**
|
||||
- [ ] **3.** `batchcmd.go`. **4.** `main.go`. **5.** README and skill.
|
||||
- [ ] **6.** `gofmt -w`; `go test -race -count=3 ./cmd/sift/` → `ok`. **7.** `make gate`.
|
||||
- [ ] **8.** Row `p6/09-cli-docs`; commit.
|
||||
|
||||
```sh
|
||||
git add cmd/sift README.md skill/SKILL.md docs/implementer-log.md
|
||||
git commit -F .state/commit-msg.txt
|
||||
```
|
||||
|
||||
## Done when
|
||||
|
||||
- The given test and every existing cmd test pass under `-race -count=3`; `make gate` ok; the
|
||||
given file is byte-identical.
|
||||
@@ -0,0 +1,87 @@
|
||||
# P6 implementation plan: asynchronous batch fetch
|
||||
|
||||
> **For the implementing model:** do not work from this file. The owner gives you one task file
|
||||
> at a time. This file is the index for the owner and the reviewer.
|
||||
|
||||
**Goal:** an agent that has collected a series of targets (URLs) and wants them fetched without
|
||||
holding a request open per target submits them as one batch, gets a job ID back immediately, and
|
||||
polls for per-target results later. A batch is **deliberate and slow** — serial, one target at a
|
||||
time — so it is the tool for research threads, not for interactive fetches.
|
||||
|
||||
Decisions (settled with the owner, 2026-09-26; revised after review the same evening, see
|
||||
`docs/P6.md` §14):
|
||||
|
||||
- **Serial per job, few workers.** One worker by default (at most 2) drains all live jobs,
|
||||
round-robin, one target at a time. The batch therefore holds at most 1–2 of the 8 global fetch
|
||||
slots, so interactive requests are never starved. A target refused busy (fetch gate or PDF
|
||||
pool) is retried after a pause, for up to 10 minutes, then failed.
|
||||
- **The caller travels with the job.** The runner keeps each job's `*auth.Caller` in memory and
|
||||
hands it to every fetch, so per-caller gates, the PDF pool, the digest cache class and
|
||||
restricted-caller redaction all apply. Nothing about the caller is persisted beyond its name
|
||||
and class: jobs never resume after a restart.
|
||||
- **Caps.** targets/job 32; live jobs/caller 1; live jobs total 8; body 256 KiB. All instrumented
|
||||
(`sift_batch_*`).
|
||||
- **Idempotency in v1.** The same (caller, target list) returns the existing job while it is live or
|
||||
has results (`queued`/`running`/`done`/`partial`) and its lease holds. A failed or cancelled job
|
||||
is not reused.
|
||||
- **Storage.** `batch_job` + `batch_target` tables in the existing `dcache` SQLite, with its own
|
||||
byte budget (64 MiB), outside the content cache's LRU.
|
||||
- **Retention / TTL.** One sliding lease per job, bumped on submit, on every target write, and on
|
||||
every owner read (the API calls `Touch`; store reads never write). GC deletes a job when its
|
||||
lease is older than the TTL (24 h).
|
||||
- **Reads.** `GET /v1/batch/{id}` is a summary (statuses, errors). A result is read one at a time
|
||||
(`GET /v1/batch/{id}/{idx}`), or all at once with `?results=1`. 32 inline results can be 6 MB.
|
||||
- **Restart.** `queued`/`running` jobs are marked `failed` (`restart`) on startup; never resumed.
|
||||
- **Restricted callers.** Their fetches go through the same `fetchTarget` with the real caller,
|
||||
so the existing redaction applies. No caller can see another's job (404, never 403). Batch
|
||||
log lines never carry a target URL or hostname.
|
||||
|
||||
| Task | Package | Given |
|
||||
|---|---|---|
|
||||
| 01-batchstore | `internal/dcache` | `dcache/batch_test.go` |
|
||||
| 02-runner-core | `internal/batch` (types, Submit, Status, Touch, Cancel; no workers) | `batch/runner_core_test.go` |
|
||||
| 03-runner-workers | `internal/batch` (workers, retry, cancel in flight, hooks) | `batch/runner_test.go` |
|
||||
| 04-metrics | `internal/metrics` | `metrics/batch_test.go` |
|
||||
| 05-fetchtarget | `internal/api` (the one refactor of working code) | `api/fetchtarget_test.go` |
|
||||
| 06-batch-wiring | `internal/api` | `api/batchrun_test.go` |
|
||||
| 07-api-endpoints | `internal/api` | `api/batch_test.go` |
|
||||
| 08-serve-config | `internal/config`, `cmd/sift/serve.go` | `config/batch_test.go` |
|
||||
| 09-cli-docs | `cmd/sift`, README, skill | `cmd_sift/batch_test.go` |
|
||||
|
||||
`guard.PreCheck` and its test are on master (owner commit): `AGENTS.md` forbids the implementer
|
||||
to edit `internal/guard/`, the security boundary.
|
||||
|
||||
Given files live under `_files/<package>/` (`cmd_sift` is `cmd/sift`); each task says where to
|
||||
copy its file. Every given test was compiled against a panic-only skeleton of the whole plan and
|
||||
fails there on `panic: todo`, not on a compile error.
|
||||
|
||||
**How this plan was made:** design worked through with the owner from `docs/PLAN.md` §13 (future
|
||||
work) on top of the P5 pool (`internal/pdfpool`), cache (`internal/dcache`), and gate
|
||||
(`global`/`per` in `internal/api`) infrastructure. Acceptance tests first: the given tests are
|
||||
authored against this spec and compiled against a panic-only skeleton so they fail for the
|
||||
intended reasons before each task is implemented.
|
||||
|
||||
## Changes during the run
|
||||
|
||||
- **Task 02 (runner), first attempt, 20:39–21:08. Model and task size, in about equal parts.**
|
||||
- The session wrote the whole runner out as ~700 lines of reasoning in its reply, never in a
|
||||
file. OpenCode compacted the context, and the model then emitted a malformed tool call
|
||||
(`</tool_call>`) and ended with no commit and no file.
|
||||
- It had also looked for the given test at a wrong path. The task said `_files/…` relative
|
||||
to the plan; that is my ambiguity.
|
||||
- The owner changes:
|
||||
- task 02 split into **02-runner-core** (no goroutines) and **03-runner-workers**, with the
|
||||
given test split to match (`runner_core_test.go` + `runner_test.go`); later tasks
|
||||
renumbered 04–09;
|
||||
- every copy path is spelled from the repo root (`docs/plans/p6/_files/…`);
|
||||
- `AGENTS.md` rule 6: write code into files early, in small steps;
|
||||
- `guard.PreCheck` moved to an owner commit on master. The old task 05 (now 06) told the
|
||||
implementer to edit `internal/guard/`, which `AGENTS.md` forbids: my contradiction, found
|
||||
while fixing the above.
|
||||
- **Task 03 (workers), first attempt, 21:31–21:48.** Same mode again: the model deliberated over
|
||||
the concurrency design in one reply, which ended mid-sentence. That was probably the 32k
|
||||
output cap, and a reply without a tool call ends `opencode run`. No file was written.
|
||||
`AGENTS.md` rule 6 alone did not prevent it. The owner changes: task 03 gains a **Design**
|
||||
section that decides the concurrency (state, the worker loop, Cancel, Close), so the task is
|
||||
transcription, not design. Task 02 (core) passed first time (`e2ddb2e`).
|
||||
|
||||
@@ -0,0 +1,280 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Every test fetches only from local httptest servers or internal.test (fake DNS): no test
|
||||
// touches the network. newBatchServer, pageServer and fakeDNS come from batchrun_test.go.
|
||||
|
||||
func batchReq(t *testing.T, s *Server, method, path, key string, body any) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
var rd io.Reader
|
||||
if body != nil {
|
||||
b, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
rd = strings.NewReader(string(b))
|
||||
}
|
||||
req := httptest.NewRequest(method, path, rd)
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
if body != nil {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
s.ServeHTTP(rec, req)
|
||||
return rec
|
||||
}
|
||||
|
||||
func batchDecode(t *testing.T, rec *httptest.ResponseRecorder) map[string]any {
|
||||
t.Helper()
|
||||
var m map[string]any
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &m); err != nil {
|
||||
t.Fatalf("decode: %v (%d %s)", err, rec.Code, rec.Body.String())
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func submit(t *testing.T, s *Server, key string, targets ...map[string]any) map[string]any {
|
||||
t.Helper()
|
||||
rec := batchReq(t, s, http.MethodPost, "/v1/batch", key, map[string]any{"targets": targets})
|
||||
if rec.Code != 202 {
|
||||
t.Fatalf("submit: %d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
return batchDecode(t, rec)
|
||||
}
|
||||
|
||||
func waitTerminal(t *testing.T, s *Server, key, id string) map[string]any {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(10 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
env := batchDecode(t, batchReq(t, s, http.MethodGet, "/v1/batch/"+id, key, nil))
|
||||
switch env["status"] {
|
||||
case "done", "partial", "failed", "cancelled":
|
||||
return env
|
||||
}
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("batch never reached a terminal status")
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestBatchSubmitAndCollect(t *testing.T) {
|
||||
ts := pageServer(t, nil)
|
||||
s, keys := newBatchServer(t, "test")
|
||||
sub := submit(t, s, keys["test"], map[string]any{"url": ts.URL + "/one"}, map[string]any{"url": ts.URL + "/two"})
|
||||
id, _ := sub["job_id"].(string)
|
||||
if sub["success"] != true || id == "" || sub["reused"] != false || sub["targets"] != float64(2) || sub["status"] != "queued" {
|
||||
t.Fatalf("submit body %v", sub)
|
||||
}
|
||||
|
||||
// The status body is a summary: statuses, no results.
|
||||
env := waitTerminal(t, s, keys["test"], id)
|
||||
if env["status"] != "done" || env["targets_total"] != float64(2) || env["targets_done"] != float64(2) {
|
||||
t.Fatalf("status %v", env)
|
||||
}
|
||||
if _, err := time.Parse(time.RFC3339, fmt.Sprint(env["created_at"])); err != nil {
|
||||
t.Fatalf("created_at %v is not RFC 3339", env["created_at"])
|
||||
}
|
||||
targets := env["targets"].([]any)
|
||||
if len(targets) != 2 {
|
||||
t.Fatalf("targets %d", len(targets))
|
||||
}
|
||||
for i, x := range targets {
|
||||
tg := x.(map[string]any)
|
||||
if tg["idx"] != float64(i) || tg["status"] != "done" {
|
||||
t.Fatalf("target %d: %v", i, tg)
|
||||
}
|
||||
if _, has := tg["result"]; has {
|
||||
t.Fatalf("summary must not carry results: target %d has one", i)
|
||||
}
|
||||
}
|
||||
|
||||
// One target's result: the /v1/fetch envelope.
|
||||
rec := batchReq(t, s, http.MethodGet, "/v1/batch/"+id+"/1", keys["test"], nil)
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("GET idx: %d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
one := batchDecode(t, rec)
|
||||
res, _ := one["result"].(map[string]any)
|
||||
if one["idx"] != float64(1) || one["status"] != "done" || res == nil {
|
||||
t.Fatalf("GET idx body %v", one)
|
||||
}
|
||||
if c, _ := res["content"].(string); !strings.Contains(c, "hello batch /two") {
|
||||
t.Fatalf("result content %q", c)
|
||||
}
|
||||
if res["untrusted"] != true || res["mode"] != "clean" {
|
||||
t.Fatalf("result is not the /v1/fetch envelope: %v", res)
|
||||
}
|
||||
if src, _ := res["source"].(map[string]any); src == nil || src["final_url"] != ts.URL+"/two" {
|
||||
t.Fatalf("result source %v", res["source"])
|
||||
}
|
||||
|
||||
// ?results=1 inlines every result.
|
||||
full := batchDecode(t, batchReq(t, s, http.MethodGet, "/v1/batch/"+id+"?results=1", keys["test"], nil))
|
||||
for i, x := range full["targets"].([]any) {
|
||||
if _, has := x.(map[string]any)["result"]; !has {
|
||||
t.Fatalf("results=1: target %d has no result", i)
|
||||
}
|
||||
}
|
||||
|
||||
if got := batchReq(t, s, http.MethodGet, "/v1/batch/"+id+"/2", keys["test"], nil).Code; got != 404 {
|
||||
t.Fatalf("GET idx out of range: %d, want 404", got)
|
||||
}
|
||||
if got := batchReq(t, s, http.MethodGet, "/v1/batch/"+id+"/x", keys["test"], nil).Code; got != 404 {
|
||||
t.Fatalf("GET idx not a number: %d, want 404", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchIdempotentResubmit(t *testing.T) {
|
||||
ts := pageServer(t, nil)
|
||||
s, keys := newBatchServer(t, "test")
|
||||
tg := map[string]any{"url": ts.URL, "mode": "clean"}
|
||||
a := submit(t, s, keys["test"], tg)
|
||||
b := submit(t, s, keys["test"], tg)
|
||||
if a["job_id"] != b["job_id"] || b["reused"] != true {
|
||||
t.Fatalf("resubmit: %v then %v", a, b)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchRefusedUpFront(t *testing.T) {
|
||||
s, keys := newBatchServer(t, "test")
|
||||
cases := []struct {
|
||||
name string
|
||||
body any
|
||||
code int
|
||||
want string
|
||||
}{
|
||||
{"guard pre-check (port)", map[string]any{"targets": []any{map[string]any{"url": "http://127.0.0.1/"}, map[string]any{"url": "http://example.com:8080/"}}}, 422, "blocked"},
|
||||
{"guard pre-check (private literal)", map[string]any{"targets": []any{map[string]any{"url": "http://10.0.0.5/"}}}, 422, "blocked"},
|
||||
{"no targets", map[string]any{"targets": []any{}}, 400, ""},
|
||||
{"not json", "nope", 400, ""},
|
||||
{"empty url", map[string]any{"targets": []any{map[string]any{"url": ""}}}, 400, ""},
|
||||
{"bad mode", map[string]any{"targets": []any{map[string]any{"url": "http://127.0.0.1/", "mode": "shout"}}}, 400, ""},
|
||||
}
|
||||
for _, c := range cases {
|
||||
rec := batchReq(t, s, http.MethodPost, "/v1/batch", keys["test"], c.body)
|
||||
if rec.Code != c.code || !strings.Contains(rec.Body.String(), c.want) {
|
||||
t.Errorf("%s: %d %s, want %d containing %q", c.name, rec.Code, rec.Body.String(), c.code, c.want)
|
||||
}
|
||||
}
|
||||
// The guard refusal names the offending index.
|
||||
rec := batchReq(t, s, http.MethodPost, "/v1/batch", keys["test"], cases[0].body)
|
||||
if !strings.Contains(rec.Body.String(), "1") {
|
||||
t.Errorf("guard refusal does not name target 1: %s", rec.Body.String())
|
||||
}
|
||||
|
||||
var many []any
|
||||
for i := 0; i < 33; i++ {
|
||||
many = append(many, map[string]any{"url": fmt.Sprintf("http://127.0.0.1/%d", i)})
|
||||
}
|
||||
if rec := batchReq(t, s, http.MethodPost, "/v1/batch", keys["test"], map[string]any{"targets": many}); rec.Code != 422 {
|
||||
t.Errorf("33 targets: %d, want 422", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// 32 targets, each with a 500-rune question, fit the batch body cap even though they are well
|
||||
// over the 64 KiB cap of the other endpoints.
|
||||
func TestBatchBodyCap(t *testing.T) {
|
||||
hold := make(chan struct{})
|
||||
defer close(hold)
|
||||
ts := pageServer(t, hold)
|
||||
s, keys := newBatchServer(t, "test")
|
||||
q := strings.Repeat("q", 500)
|
||||
var many []any
|
||||
for i := 0; i < 32; i++ {
|
||||
many = append(many, map[string]any{"url": fmt.Sprintf("%s/%d?pad=%s", ts.URL, i, strings.Repeat("p", 1500)), "question": q})
|
||||
}
|
||||
if rec := batchReq(t, s, http.MethodPost, "/v1/batch", keys["test"], map[string]any{"targets": many}); rec.Code != 202 {
|
||||
t.Fatalf("32 targets with questions (~%d KiB): %d %s", 32*2100/1024, rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchPartialFailure(t *testing.T) {
|
||||
ts := pageServer(t, nil)
|
||||
s, keys := newBatchServer(t, "test")
|
||||
// Target 1 is a name that resolves to a private address. Names are not resolved at submit,
|
||||
// so it passes the pre-check; the fetch-time guard blocks it: a per-target failure.
|
||||
sub := submit(t, s, keys["test"], map[string]any{"url": ts.URL}, map[string]any{"url": "http://internal.test/"})
|
||||
env := waitTerminal(t, s, keys["test"], sub["job_id"].(string))
|
||||
if env["status"] != "partial" {
|
||||
t.Fatalf("status %v", env["status"])
|
||||
}
|
||||
second := env["targets"].([]any)[1].(map[string]any)
|
||||
if second["status"] != "failed" {
|
||||
t.Fatalf("target 1: %v", second)
|
||||
}
|
||||
if e, _ := second["error"].(string); !strings.Contains(e, "blocked") {
|
||||
t.Fatalf("target 1 error %q", e)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchOwnership404(t *testing.T) {
|
||||
hold := make(chan struct{})
|
||||
defer close(hold)
|
||||
ts := pageServer(t, hold)
|
||||
s, keys := newBatchServer(t, "test", "other")
|
||||
id := submit(t, s, keys["test"], map[string]any{"url": ts.URL})["job_id"].(string)
|
||||
for _, p := range []string{"/v1/batch/" + id, "/v1/batch/" + id + "/0"} {
|
||||
if got := batchReq(t, s, http.MethodGet, p, keys["other"], nil).Code; got != 404 {
|
||||
t.Fatalf("other caller GET %s: %d (must not reveal the job)", p, got)
|
||||
}
|
||||
}
|
||||
if got := batchReq(t, s, http.MethodDelete, "/v1/batch/"+id, keys["other"], nil).Code; got != 404 {
|
||||
t.Fatalf("other caller DELETE: %d", got)
|
||||
}
|
||||
if got := batchReq(t, s, http.MethodGet, "/v1/batch/"+id, keys["test"], nil).Code; got != 200 {
|
||||
t.Fatalf("owner GET: %d", got)
|
||||
}
|
||||
if got := batchReq(t, s, http.MethodGet, "/v1/batch/no-such-job", keys["test"], nil).Code; got != 404 {
|
||||
t.Fatalf("unknown job: %d", got)
|
||||
}
|
||||
if got := batchReq(t, s, http.MethodGet, "/v1/batch/"+id, "not-a-key", nil).Code; got != 401 {
|
||||
t.Fatalf("bad key: %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchCancel(t *testing.T) {
|
||||
hold := make(chan struct{})
|
||||
ts := pageServer(t, hold)
|
||||
s, keys := newBatchServer(t, "test")
|
||||
id := submit(t, s, keys["test"], map[string]any{"url": ts.URL + "/a"}, map[string]any{"url": ts.URL + "/b"})["job_id"].(string)
|
||||
rec := batchReq(t, s, http.MethodDelete, "/v1/batch/"+id, keys["test"], nil)
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("cancel: %d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if b := batchDecode(t, rec); b["success"] != true || b["status"] != "cancelled" {
|
||||
t.Fatalf("cancel body %v", b)
|
||||
}
|
||||
close(hold)
|
||||
env := waitTerminal(t, s, keys["test"], id)
|
||||
if env["status"] != "cancelled" {
|
||||
t.Fatalf("status %v", env["status"])
|
||||
}
|
||||
if tg := env["targets"].([]any)[1].(map[string]any); tg["status"] != "pending" {
|
||||
t.Fatalf("target 1 after cancel: %v", tg)
|
||||
}
|
||||
if got := batchReq(t, s, http.MethodDelete, "/v1/batch/"+id, keys["test"], nil).Code; got != 409 {
|
||||
t.Fatalf("cancelling a finished job: %d, want 409", got)
|
||||
}
|
||||
}
|
||||
|
||||
// The request metrics never get a job id as a label value.
|
||||
func TestBatchMetricsPathLabel(t *testing.T) {
|
||||
ts := pageServer(t, nil)
|
||||
s, keys := newBatchServer(t, "test")
|
||||
id := submit(t, s, keys["test"], map[string]any{"url": ts.URL})["job_id"].(string)
|
||||
waitTerminal(t, s, keys["test"], id)
|
||||
batchReq(t, s, http.MethodGet, "/v1/batch/"+id+"/0", keys["test"], nil)
|
||||
if strings.Contains(promText(t, s), id) {
|
||||
t.Fatal("a job id appears in /metrics")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
package api
|
||||
|
||||
// P6 task 06: the runner wired to fetchTarget and the metrics, with no HTTP yet. The helpers
|
||||
// here (newBatchServer, pageServer, fakeDNS) are shared with task 07's batch_test.go.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/netip"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/auth"
|
||||
"git.wntrmute.dev/kyle/sift/internal/batch"
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
"git.wntrmute.dev/kyle/sift/internal/fetch"
|
||||
"git.wntrmute.dev/kyle/sift/internal/guard"
|
||||
)
|
||||
|
||||
var loopback = []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}
|
||||
|
||||
// fakeDNS resolves internal.test to a private address; every other name fails to resolve.
|
||||
func fakeDNS(_ context.Context, host string) ([]netip.Addr, error) {
|
||||
if host == "internal.test" {
|
||||
return []netip.Addr{netip.MustParseAddr("10.0.0.5")}, nil
|
||||
}
|
||||
return nil, fmt.Errorf("no such host %s", host)
|
||||
}
|
||||
|
||||
// newBatchServer builds a Server with the batch runner wired over an in-memory dcache through
|
||||
// s.NewBatchRunner, exactly as serve does. callers are the names to provision (each gets a key).
|
||||
func newBatchServer(t *testing.T, callers ...string) (*Server, map[string]string) {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
keysPath := filepath.Join(dir, "keys.toml")
|
||||
keys := map[string]string{}
|
||||
for _, name := range callers {
|
||||
out := filepath.Join(dir, name)
|
||||
if err := auth.Add(keysPath, name, "hermes-fleet", out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
b, err := os.ReadFile(out)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
keys[name] = strings.TrimSpace(string(b))
|
||||
}
|
||||
store, err := auth.Load(keysPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
g := &guard.Guard{AllowPrefixes: loopback, LookupIP: fakeDNS}
|
||||
f := &fetch.Fetcher{Guard: g, UA: "sift-test"}
|
||||
log := slog.New(slog.NewTextHandler(io.Discard, nil))
|
||||
s := New("test", g, f, store, log)
|
||||
db, err := dcache.Open("", 1<<20, log)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
s.Batch = s.NewBatchRunner(db, batch.Config{})
|
||||
t.Cleanup(s.Batch.Close) // before db.Close
|
||||
return s, keys
|
||||
}
|
||||
|
||||
// pageServer serves a small HTML page at every path; hold, when non-nil, blocks each request
|
||||
// until it is closed.
|
||||
func pageServer(t *testing.T, hold chan struct{}) *httptest.Server {
|
||||
t.Helper()
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if hold != nil {
|
||||
select {
|
||||
case <-hold:
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
w.Header().Set("Content-Type", "text/html; charset=utf-8")
|
||||
fmt.Fprintf(w, `<html><head><title>T</title></head><body><article><p>hello batch %s</p></article></body></html>`, r.URL.Path)
|
||||
}))
|
||||
t.Cleanup(ts.Close)
|
||||
return ts
|
||||
}
|
||||
|
||||
func waitJob(t *testing.T, r *batch.Runner, id string) batch.Job {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(10 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
j, _, ok := r.Status(id)
|
||||
if ok && (j.Status == "done" || j.Status == "partial" || j.Status == "failed" || j.Status == "cancelled") {
|
||||
return j
|
||||
}
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("job never finished")
|
||||
return batch.Job{}
|
||||
}
|
||||
|
||||
func promText(t *testing.T, s *Server) string {
|
||||
t.Helper()
|
||||
req := httptest.NewRequest(http.MethodGet, "/metrics", nil)
|
||||
req.RemoteAddr, req.Host = "127.0.0.1:9", "127.0.0.1:8890"
|
||||
rec := httptest.NewRecorder()
|
||||
s.ServeHTTP(rec, req)
|
||||
return rec.Body.String()
|
||||
}
|
||||
|
||||
// The runner fetches through fetchTarget as the submitting caller; a name that resolves to a
|
||||
// private address is a per-target "blocked" failure; the hooks reach the metrics.
|
||||
func TestBatchRunnerWired(t *testing.T) {
|
||||
ts := pageServer(t, nil)
|
||||
s, _ := newBatchServer(t, "test")
|
||||
c := &auth.Caller{Name: "test", Workspace: "hermes-fleet"}
|
||||
id, reused, err := s.Batch.Submit(context.Background(), c, dcache.ClassFleet, []batch.Target{
|
||||
{URL: ts.URL + "/one"}, {URL: "http://internal.test/"},
|
||||
})
|
||||
if err != nil || reused {
|
||||
t.Fatalf("submit: %v %v", err, reused)
|
||||
}
|
||||
if j := waitJob(t, s.Batch, id); j.Status != "partial" {
|
||||
t.Fatalf("job %+v", j)
|
||||
}
|
||||
_, ts2, _ := s.Batch.Status(id)
|
||||
if ts2[0].Status != "done" || ts2[0].Result == nil || !strings.Contains(ts2[0].Result.Content, "hello batch /one") {
|
||||
t.Fatalf("target 0: %+v", ts2[0])
|
||||
}
|
||||
if ts2[0].Result.Meta["mode"] != "clean" || ts2[0].Result.Meta["untrusted"] != true {
|
||||
t.Fatalf("target 0 meta %v: not the fetchTarget envelope", ts2[0].Result.Meta)
|
||||
}
|
||||
if ts2[1].Status != "failed" || !strings.Contains(ts2[1].Error, "blocked") {
|
||||
t.Fatalf("target 1: %+v", ts2[1])
|
||||
}
|
||||
out := promText(t, s)
|
||||
for _, want := range []string{
|
||||
`sift_batch_jobs_total{outcome="created"} 1`,
|
||||
`sift_batch_jobs_total{outcome="partial"} 1`,
|
||||
`sift_batch_targets_total{outcome="ok"} 1`,
|
||||
`sift_batch_targets_total{outcome="blocked"} 1`,
|
||||
`sift_batch_result_bytes_count 1`,
|
||||
`sift_batch_queue_depth 0`,
|
||||
} {
|
||||
if !strings.Contains(out, want) {
|
||||
t.Errorf("metrics lack %q", want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The adapter turns fetch errors into what the runner and the metrics need.
|
||||
func TestBatchOutcomeClasses(t *testing.T) {
|
||||
cases := []struct {
|
||||
err error
|
||||
want string
|
||||
}{
|
||||
{nil, "ok"},
|
||||
{fmt.Errorf("busy: %w", batch.ErrRetryLater), "busy"},
|
||||
{&batch.TargetError{Status: 403, Msg: "blocked: private address"}, "blocked"},
|
||||
{&batch.TargetError{Status: 429, Msg: "busy"}, "busy"},
|
||||
{&batch.TargetError{Status: 504, Msg: "timeout"}, "timeout"},
|
||||
{&batch.TargetError{Status: 415, Msg: "unsupported content type"}, "unsupported"},
|
||||
{&batch.TargetError{Status: 422, Msg: "document too long for question (6001 chunks); use grep or pages"}, "too_long"},
|
||||
{&batch.TargetError{Status: 422, Msg: "unreadable PDF"}, "upstream_4xx"},
|
||||
{&batch.TargetError{Status: 404, Msg: "not found"}, "upstream_4xx"},
|
||||
{&batch.TargetError{Status: 500, Msg: "internal error"}, "upstream_5xx"},
|
||||
{errors.New("something else"), "upstream_5xx"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := batchOutcome(c.err); got != c.want {
|
||||
t.Errorf("batchOutcome(%v) = %q, want %q", c.err, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// batchFetch maps a busy fetch to ErrRetryLater (the runner retries it) and every other
|
||||
// failure to a TargetError with the /v1/fetch status.
|
||||
func TestBatchFetchErrors(t *testing.T) {
|
||||
s, _ := newBatchServer(t, "test")
|
||||
c := &auth.Caller{Name: "test", Workspace: "hermes-fleet"}
|
||||
_, err := s.batchFetch(context.Background(), c, batch.Target{URL: "http://internal.test/"})
|
||||
var te *batch.TargetError
|
||||
if !errors.As(err, &te) || te.Status != 403 || !strings.Contains(te.Msg, "blocked") {
|
||||
t.Fatalf("private name: %v, want a TargetError 403 blocked", err)
|
||||
}
|
||||
// Fill the global fetch gate so the next fetch is refused busy.
|
||||
for i := 0; i < 8; i++ {
|
||||
if err := s.acquire(fmt.Sprintf("filler%d", i), "http://filler.test/"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
ts := pageServer(t, nil)
|
||||
_, err = s.batchFetch(context.Background(), c, batch.Target{URL: ts.URL})
|
||||
if !errors.Is(err, batch.ErrRetryLater) {
|
||||
t.Fatalf("busy gate: %v, want an error wrapping batch.ErrRetryLater", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
package api
|
||||
|
||||
// P6 task 05: v1fetch's per-target pipeline is extracted into s.fetchTarget, which the batch
|
||||
// runner also calls. These tests pin that the two agree; every existing v1fetch test must also
|
||||
// pass unchanged.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/netip"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/auth"
|
||||
"git.wntrmute.dev/kyle/sift/internal/batch"
|
||||
"git.wntrmute.dev/kyle/sift/internal/fetch"
|
||||
"git.wntrmute.dev/kyle/sift/internal/guard"
|
||||
)
|
||||
|
||||
// newTargetServer is setupAPI with loopback allowed; its key belongs to caller "test".
|
||||
func newTargetServer(t *testing.T) (*Server, string) {
|
||||
t.Helper()
|
||||
return setupAPI(t, &guard.Guard{AllowPrefixes: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}})
|
||||
}
|
||||
|
||||
func TestFetchTargetMatchesV1Fetch(t *testing.T) {
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/html; charset=utf-8")
|
||||
_, _ = w.Write([]byte(`<html><head><title>Same</title></head><body><article><h1>Same</h1><p>one pipeline, two callers</p></article></body></html>`))
|
||||
}))
|
||||
defer ts.Close()
|
||||
s, key := newTargetServer(t)
|
||||
c := &auth.Caller{Name: "test", Workspace: "hermes-fleet"}
|
||||
|
||||
res, err := s.fetchTarget(context.Background(), c, batch.Target{URL: ts.URL})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(res.Content, "one pipeline, two callers") {
|
||||
t.Fatalf("content %q", res.Content)
|
||||
}
|
||||
if _, has := res.Meta["content"]; has {
|
||||
t.Fatal("Meta must not repeat the content")
|
||||
}
|
||||
for _, k := range []string{"mode", "untrusted", "source", "warnings", "sizes", "timings"} {
|
||||
if _, has := res.Meta[k]; !has {
|
||||
t.Errorf("Meta lacks %q (it carries every envelope field except content)", k)
|
||||
}
|
||||
}
|
||||
if res.Meta["mode"] != "clean" || res.Meta["untrusted"] != true {
|
||||
t.Fatalf("meta %v", res.Meta)
|
||||
}
|
||||
|
||||
// The same URL through /v1/fetch: identical content and mode, and the envelope has exactly
|
||||
// the fields of Meta plus content.
|
||||
rec := postJSON(t, s, "/v1/fetch", key, map[string]any{"url": ts.URL})
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("v1fetch %d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var env map[string]any
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &env); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if env["content"] != res.Content || env["mode"] != res.Meta["mode"] {
|
||||
t.Fatalf("v1fetch and fetchTarget disagree:\n%v\n%+v", env, res)
|
||||
}
|
||||
if len(env) != len(res.Meta)+1 {
|
||||
t.Fatalf("envelope keys %d, Meta keys %d: want Meta + content", len(env), len(res.Meta))
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchTargetErrors(t *testing.T) {
|
||||
s, _ := newTargetServer(t)
|
||||
c := &auth.Caller{Name: "test", Workspace: "hermes-fleet"}
|
||||
cases := []struct {
|
||||
name string
|
||||
target batch.Target
|
||||
status int
|
||||
msg string
|
||||
}{
|
||||
{"blocked", batch.Target{URL: "http://10.0.0.5/"}, 403, "blocked"},
|
||||
{"bad mode", batch.Target{URL: "http://127.0.0.1/", Options: batch.Options{Mode: "shout"}}, 400, ""},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
_, err := s.fetchTarget(context.Background(), c, tc.target)
|
||||
var fe *fetch.Error
|
||||
if !errors.As(err, &fe) || fe.Status != tc.status || !strings.Contains(fe.Msg, tc.msg) {
|
||||
t.Errorf("%s: err %v, want a *fetch.Error %d containing %q", tc.name, err, tc.status, tc.msg)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,182 @@
|
||||
package batch
|
||||
|
||||
// P6 task 02: the runner's core — types, Submit (idempotency, caps, storage), Status, Touch and
|
||||
// Cancel — with no workers yet: every job stays queued. The helpers here are shared with task
|
||||
// 03's runner_test.go.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"log/slog"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/auth"
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
)
|
||||
|
||||
var (
|
||||
alice = &auth.Caller{Name: "alice", Workspace: "hermes-fleet"}
|
||||
bob = &auth.Caller{Name: "bob", Workspace: "hermes-fleet"}
|
||||
carol = &auth.Caller{Name: "carol", Workspace: "hermes-fleet"}
|
||||
)
|
||||
|
||||
func newRunner(t *testing.T, fetch FetchFunc, cfg Config) *Runner {
|
||||
t.Helper()
|
||||
db, err := dcache.Open("", 1<<20, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
r := New(db, fetch, cfg, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
t.Cleanup(r.Close) // runs before db.Close (cleanups are LIFO)
|
||||
return r
|
||||
}
|
||||
|
||||
func waitStatus(t *testing.T, r *Runner, id, want string) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
j, _, ok := r.Status(id)
|
||||
if ok && j.Status == want {
|
||||
return
|
||||
}
|
||||
time.Sleep(2 * time.Millisecond)
|
||||
}
|
||||
j, _, _ := r.Status(id)
|
||||
t.Fatalf("job %s never reached %q (now %q)", id, want, j.Status)
|
||||
}
|
||||
|
||||
func urls(ss ...string) []Target {
|
||||
var out []Target
|
||||
for _, s := range ss {
|
||||
out = append(out, Target{URL: "https://" + s + ".example/"})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Submit stores the job queued with its targets pending, before it returns.
|
||||
func TestSubmitStoresQueued(t *testing.T) {
|
||||
r := newRunner(t, func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
return Result{Content: "ok"}, nil
|
||||
}, Config{Workers: 1})
|
||||
id, reused, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
if err != nil || reused || len(id) != 32 {
|
||||
t.Fatalf("submit id=%q reused=%v err=%v, want a 32-hex-char id", id, reused, err)
|
||||
}
|
||||
j, ts, ok := r.Status(id)
|
||||
if !ok || j.ID != id || j.Caller != "alice" || j.TargetsTotal != 2 || j.CreatedAt == 0 {
|
||||
t.Fatalf("job %+v ok=%v", j, ok)
|
||||
}
|
||||
if j.Status != "queued" && j.Status != "running" && j.Status != "done" {
|
||||
t.Fatalf("status %q", j.Status)
|
||||
}
|
||||
if len(ts) != 2 || ts[0].Idx != 0 || ts[0].URL != "https://a.example/" || ts[1].URL != "https://b.example/" {
|
||||
t.Fatalf("targets %+v", ts)
|
||||
}
|
||||
}
|
||||
|
||||
// Cancelling a job no worker has reached: it is cancelled at once, its targets stay pending,
|
||||
// and the same list then starts a new job.
|
||||
func TestCancelQueued(t *testing.T) {
|
||||
release := make(chan struct{})
|
||||
defer close(release)
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
<-release
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{Workers: 1})
|
||||
busy, _, err := r.Submit(context.Background(), bob, dcache.ClassFleet, urls("x")) // occupies the only worker once there is one
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_ = busy
|
||||
id, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !r.Cancel(id) {
|
||||
t.Fatal("cancel returned false")
|
||||
}
|
||||
j, ts, _ := r.Status(id)
|
||||
if j.Status != "cancelled" || ts[0].Status != "pending" || ts[1].Status != "pending" {
|
||||
t.Fatalf("after cancel: %+v %+v", j, ts)
|
||||
}
|
||||
if r.Cancel(id) || r.Cancel("no-such-job") {
|
||||
t.Fatal("cancel of a finished or unknown job must return false")
|
||||
}
|
||||
id2, reused, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
if err != nil || reused || id2 == id {
|
||||
t.Fatalf("resubmit after cancel id=%q reused=%v err=%v", id2, reused, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Touch keeps a job alive; Status never does.
|
||||
func TestTouchAndStatusLease(t *testing.T) {
|
||||
db, err := dcache.Open("", 1<<20, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
release := make(chan struct{})
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
<-release
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := New(db, fetch, Config{Workers: 1}, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
t.Cleanup(r.Close)
|
||||
defer close(release) // runs before the cleanups
|
||||
id, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond) // let a worker (task 03) take the target; it then blocks
|
||||
before, _ := db.BatchJob(id)
|
||||
time.Sleep(1100 * time.Millisecond) // the lease is in whole seconds
|
||||
r.Status(id)
|
||||
if j, _ := db.BatchJob(id); j.Lease != before.Lease {
|
||||
t.Fatalf("Status moved the lease %d → %d; a read must not write", before.Lease, j.Lease)
|
||||
}
|
||||
r.Touch(id)
|
||||
if j, _ := db.BatchJob(id); j.Lease <= before.Lease {
|
||||
t.Fatalf("Touch did not move the lease: %d → %d", before.Lease, j.Lease)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCaps(t *testing.T) {
|
||||
release := make(chan struct{})
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
<-release
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{TargetsPerJob: 2, JobsTotal: 2})
|
||||
defer close(release)
|
||||
if _, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b", "c")); !errors.Is(err, ErrTooManyTargets) {
|
||||
t.Fatalf("3 targets over a cap of 2: want ErrTooManyTargets, got %v", err)
|
||||
}
|
||||
if _, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, nil); !errors.Is(err, ErrNoTargets) {
|
||||
t.Fatalf("no targets: want ErrNoTargets, got %v", err)
|
||||
}
|
||||
if _, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// alice's job is live (queued or running): a second, different list is refused.
|
||||
if _, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("z")); !errors.Is(err, ErrCallerBusy) {
|
||||
t.Fatalf("want ErrCallerBusy, got %v", err)
|
||||
}
|
||||
if _, _, err := r.Submit(context.Background(), bob, dcache.ClassFleet, urls("b")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Two live jobs in total: a third caller is refused.
|
||||
if _, _, err := r.Submit(context.Background(), carol, dcache.ClassFleet, urls("c")); !errors.Is(err, ErrBusy) {
|
||||
t.Fatalf("want ErrBusy, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStatusUnknown(t *testing.T) {
|
||||
r := newRunner(t, func(context.Context, *auth.Caller, Target) (Result, error) { return Result{}, nil }, Config{})
|
||||
if _, _, ok := r.Status("no-such-job"); ok {
|
||||
t.Fatal("unknown id must be ok=false")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,373 @@
|
||||
package batch
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/auth"
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
)
|
||||
|
||||
// Targets run one at a time, in order, and each fetch is given the submitting caller.
|
||||
func TestSerialOrderAndCaller(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
var calls, names []string
|
||||
inflight, peak := 0, 0
|
||||
fetch := func(_ context.Context, c *auth.Caller, tgt Target) (Result, error) {
|
||||
mu.Lock()
|
||||
inflight++
|
||||
if inflight > peak {
|
||||
peak = inflight
|
||||
}
|
||||
calls = append(calls, tgt.URL)
|
||||
names = append(names, c.Name)
|
||||
mu.Unlock()
|
||||
time.Sleep(time.Millisecond)
|
||||
mu.Lock()
|
||||
inflight--
|
||||
mu.Unlock()
|
||||
return Result{Content: "ok " + tgt.URL}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{})
|
||||
id, reused, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b", "c"))
|
||||
if err != nil || reused {
|
||||
t.Fatalf("submit err=%v reused=%v", err, reused)
|
||||
}
|
||||
waitStatus(t, r, id, "done")
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if strings.Join(calls, " ") != "https://a.example/ https://b.example/ https://c.example/" {
|
||||
t.Fatalf("order %v", calls)
|
||||
}
|
||||
if peak != 1 {
|
||||
t.Fatalf("peak inflight %d, want 1 (serial)", peak)
|
||||
}
|
||||
for _, n := range names {
|
||||
if n != "alice" {
|
||||
t.Fatalf("fetch saw caller %q, want alice", n)
|
||||
}
|
||||
}
|
||||
j, _, _ := r.Status(id)
|
||||
if j.Caller != "alice" || j.TargetsTotal != 3 || j.TargetsDone != 3 || j.CreatedAt == 0 {
|
||||
t.Fatalf("job %+v", j)
|
||||
}
|
||||
}
|
||||
|
||||
// With one worker, two live jobs take turns target by target.
|
||||
func TestRoundRobin(t *testing.T) {
|
||||
first := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
var once sync.Once
|
||||
var mu sync.Mutex
|
||||
var calls []string
|
||||
fetch := func(_ context.Context, _ *auth.Caller, tgt Target) (Result, error) {
|
||||
mu.Lock()
|
||||
calls = append(calls, strings.TrimSuffix(strings.TrimPrefix(tgt.URL, "https://"), ".example/"))
|
||||
mu.Unlock()
|
||||
once.Do(func() { close(first); <-release })
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{Workers: 1})
|
||||
a, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a1", "a2", "a3"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
<-first // a1 is in flight
|
||||
b, _, err := r.Submit(context.Background(), bob, dcache.ClassFleet, urls("b1", "b2", "b3"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
close(release)
|
||||
waitStatus(t, r, a, "done")
|
||||
waitStatus(t, r, b, "done")
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if got := strings.Join(calls, " "); got != "a1 b1 a2 b2 a3 b3" {
|
||||
t.Fatalf("order %q, want round-robin a1 b1 a2 b2 a3 b3", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFailureIsolation(t *testing.T) {
|
||||
fetch := func(_ context.Context, _ *auth.Caller, tgt Target) (Result, error) {
|
||||
if tgt.URL == "https://b.example/" {
|
||||
return Result{}, &TargetError{Msg: "upstream 500", Hint: "browser"}
|
||||
}
|
||||
return Result{Content: "ok " + tgt.URL}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{})
|
||||
id, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b", "c"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
waitStatus(t, r, id, "partial")
|
||||
j, ts, ok := r.Status(id)
|
||||
if !ok || j.TargetsDone != 3 {
|
||||
t.Fatalf("job %+v ok=%v", j, ok)
|
||||
}
|
||||
if ts[0].Status != "done" || ts[1].Status != "failed" || ts[2].Status != "done" {
|
||||
t.Fatalf("statuses %+v", ts)
|
||||
}
|
||||
if ts[1].Error != "upstream 500" || ts[1].Hint != "browser" {
|
||||
t.Fatalf("failed target error=%q hint=%q", ts[1].Error, ts[1].Hint)
|
||||
}
|
||||
if ts[0].Result == nil || ts[0].Result.Content != "ok https://a.example/" {
|
||||
t.Fatalf("result %+v", ts[0].Result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAllFailedIsFailed(t *testing.T) {
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) { return Result{}, errors.New("nope") }
|
||||
r := newRunner(t, fetch, Config{})
|
||||
id, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
waitStatus(t, r, id, "failed")
|
||||
_, ts, _ := r.Status(id)
|
||||
if ts[0].Error != "nope" || ts[0].Hint != "" {
|
||||
t.Fatalf("plain error stored as error=%q hint=%q", ts[0].Error, ts[0].Hint)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIdempotentResubmit(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
n := 0
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
mu.Lock()
|
||||
n++
|
||||
mu.Unlock()
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{})
|
||||
id1, reused1, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
if err != nil || reused1 {
|
||||
t.Fatalf("first submit err=%v reused=%v", err, reused1)
|
||||
}
|
||||
id2, reused2, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
if err != nil || !reused2 || id2 != id1 {
|
||||
t.Fatalf("resubmit id=%q reused=%v err=%v", id2, reused2, err)
|
||||
}
|
||||
waitStatus(t, r, id1, "done")
|
||||
// Still reused once done (within the lease), and never re-fetched.
|
||||
if id3, reused3, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b")); err != nil || !reused3 || id3 != id1 {
|
||||
t.Fatalf("resubmit after done id=%q reused=%v err=%v", id3, reused3, err)
|
||||
}
|
||||
// Another caller with the same list gets its own job.
|
||||
if id4, reused4, err := r.Submit(context.Background(), bob, dcache.ClassFleet, urls("a", "b")); err != nil || reused4 || id4 == id1 {
|
||||
t.Fatalf("other caller id=%q reused=%v err=%v", id4, reused4, err)
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if n != 4 {
|
||||
t.Fatalf("fetched %d times, want 4 (2 for alice's one job, 2 for bob's)", n)
|
||||
}
|
||||
}
|
||||
|
||||
// A cancelled job is not handed back to a resubmit: the same list starts a new job.
|
||||
func TestResubmitAfterCancelIsNew(t *testing.T) {
|
||||
release := make(chan struct{})
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
<-release
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{})
|
||||
id1, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
waitStatus(t, r, id1, "running")
|
||||
if !r.Cancel(id1) {
|
||||
t.Fatal("cancel returned false")
|
||||
}
|
||||
close(release)
|
||||
waitStatus(t, r, id1, "cancelled")
|
||||
id2, reused, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
if err != nil || reused || id2 == id1 {
|
||||
t.Fatalf("resubmit after cancel id=%q reused=%v err=%v", id2, reused, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Submit returns with the job stored "queued"; it becomes "running" when a worker takes it.
|
||||
func TestQueuedThenRunning(t *testing.T) {
|
||||
release := make(chan struct{})
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
<-release
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{Workers: 1})
|
||||
defer close(release)
|
||||
a, _, _ := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a"))
|
||||
waitStatus(t, r, a, "running")
|
||||
b, _, err := r.Submit(context.Background(), bob, dcache.ClassFleet, urls("b"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// The only worker is busy with alice's target, so bob's job waits in the queue.
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
if j, _, _ := r.Status(b); j.Status != "queued" {
|
||||
t.Fatalf("bob's job is %q while the only worker is busy, want queued", j.Status)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCancel(t *testing.T) {
|
||||
started := make(chan struct{}, 1)
|
||||
release := make(chan struct{})
|
||||
var mu sync.Mutex
|
||||
var calls []string
|
||||
fetch := func(_ context.Context, _ *auth.Caller, tgt Target) (Result, error) {
|
||||
mu.Lock()
|
||||
calls = append(calls, tgt.URL)
|
||||
mu.Unlock()
|
||||
select {
|
||||
case started <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
<-release
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{})
|
||||
id, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b", "c"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
<-started // the first target is in flight
|
||||
if !r.Cancel(id) {
|
||||
t.Fatal("cancel returned false")
|
||||
}
|
||||
close(release) // let the in-flight target finish
|
||||
waitStatus(t, r, id, "cancelled")
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
mu.Lock()
|
||||
n := len(calls)
|
||||
mu.Unlock()
|
||||
if n != 1 {
|
||||
t.Fatalf("cancelled after %d fetches, want 1", n)
|
||||
}
|
||||
j, ts, ok := r.Status(id)
|
||||
if !ok || j.Status != "cancelled" {
|
||||
t.Fatalf("job %+v: the in-flight target's write must not overwrite cancelled", j)
|
||||
}
|
||||
if ts[0].Status != "done" {
|
||||
t.Fatalf("in-flight target should be stored done: %+v", ts[0])
|
||||
}
|
||||
if ts[1].Status != "pending" || ts[2].Status != "pending" {
|
||||
t.Fatalf("later targets must stay pending: %+v", ts[1:])
|
||||
}
|
||||
if r.Cancel(id) {
|
||||
t.Fatal("cancelling a terminal job must return false")
|
||||
}
|
||||
if r.Cancel("no-such-job") {
|
||||
t.Fatal("cancelling an unknown job must return false")
|
||||
}
|
||||
}
|
||||
|
||||
// A target that reports ErrRetryLater (the fetch gate was busy) is retried, not failed.
|
||||
func TestBusyIsRetried(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
n := 0
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
n++
|
||||
if n <= 2 {
|
||||
return Result{}, fmt.Errorf("busy: %w", ErrRetryLater)
|
||||
}
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{RetryWait: time.Millisecond})
|
||||
id, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
waitStatus(t, r, id, "done")
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if n != 3 {
|
||||
t.Fatalf("fetched %d times, want 3 (two busy, then ok)", n)
|
||||
}
|
||||
}
|
||||
|
||||
// Retrying is bounded: a target busy for longer than RetryFor fails with a busy error.
|
||||
func TestBusyGivesUp(t *testing.T) {
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
return Result{}, fmt.Errorf("busy: %w", ErrRetryLater)
|
||||
}
|
||||
r := newRunner(t, fetch, Config{RetryWait: time.Millisecond, RetryFor: 30 * time.Millisecond})
|
||||
id, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
waitStatus(t, r, id, "failed")
|
||||
_, ts, _ := r.Status(id)
|
||||
if ts[0].Status != "failed" || !strings.Contains(ts[0].Error, "busy") {
|
||||
t.Fatalf("target %+v, want failed with a busy error", ts[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestPanicRecovered(t *testing.T) {
|
||||
fetch := func(_ context.Context, _ *auth.Caller, tgt Target) (Result, error) {
|
||||
if tgt.URL == "https://b.example/" {
|
||||
panic("boom")
|
||||
}
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{})
|
||||
id, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b", "c"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
waitStatus(t, r, id, "partial")
|
||||
_, ts, _ := r.Status(id)
|
||||
if ts[0].Status != "done" || ts[1].Status != "failed" || ts[2].Status != "done" {
|
||||
t.Fatalf("statuses %+v", ts)
|
||||
}
|
||||
if ts[1].Error == "" || strings.Contains(ts[1].Error, "boom") {
|
||||
t.Fatalf("panic target error %q: must be set, and must not echo the panic value", ts[1].Error)
|
||||
}
|
||||
}
|
||||
|
||||
// The hooks see every job outcome once and every stored target once.
|
||||
func TestHooks(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
jobs := map[string]int{}
|
||||
var targets, jobSecs int
|
||||
var bytes int
|
||||
h := Hooks{
|
||||
Job: func(o string) { mu.Lock(); jobs[o]++; mu.Unlock() },
|
||||
Target: func(err error, _ time.Duration, n int) {
|
||||
mu.Lock()
|
||||
targets++
|
||||
bytes += n
|
||||
mu.Unlock()
|
||||
},
|
||||
JobSeconds: func(time.Duration) { mu.Lock(); jobSecs++; mu.Unlock() },
|
||||
}
|
||||
fetch := func(_ context.Context, _ *auth.Caller, tgt Target) (Result, error) {
|
||||
if tgt.URL == "https://b.example/" {
|
||||
return Result{}, errors.New("nope")
|
||||
}
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{Hooks: h})
|
||||
id, _, _ := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
waitStatus(t, r, id, "partial")
|
||||
_, _, _ = r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b")) // reused
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if jobs["created"] != 1 || jobs["partial"] != 1 || jobs["reused"] != 1 || len(jobs) != 3 {
|
||||
t.Fatalf("job outcomes %v", jobs)
|
||||
}
|
||||
if targets != 2 || jobSecs != 1 {
|
||||
t.Fatalf("targets=%d jobSeconds=%d, want 2 and 1", targets, jobSecs)
|
||||
}
|
||||
if bytes <= 0 {
|
||||
t.Fatalf("stored result bytes %d, want > 0 for the done target", bytes)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestParseBatchArgs(t *testing.T) {
|
||||
cases := []struct {
|
||||
args []string
|
||||
verb string
|
||||
rest []string
|
||||
ok bool
|
||||
}{
|
||||
{[]string{"submit", "targets.json"}, "submit", []string{"targets.json"}, true},
|
||||
{[]string{"submit", "-"}, "submit", []string{"-"}, true},
|
||||
{[]string{"status", "job-123"}, "status", []string{"job-123"}, true},
|
||||
{[]string{"status", "--results", "job-123"}, "status", []string{"--results", "job-123"}, true},
|
||||
{[]string{"get", "job-123", "4"}, "get", []string{"job-123", "4"}, true},
|
||||
{[]string{"cancel", "job-123"}, "cancel", []string{"job-123"}, true},
|
||||
{[]string{}, "", nil, false},
|
||||
{[]string{"submit"}, "", nil, false},
|
||||
{[]string{"status"}, "", nil, false},
|
||||
{[]string{"status", "a", "b"}, "", nil, false},
|
||||
{[]string{"status", "--results"}, "", nil, false},
|
||||
{[]string{"get", "job-123"}, "", nil, false},
|
||||
{[]string{"get", "job-123", "x"}, "", nil, false},
|
||||
{[]string{"get", "job-123", "-1"}, "", nil, false},
|
||||
{[]string{"cancel"}, "", nil, false},
|
||||
{[]string{"bogus", "x"}, "", nil, false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
verb, rest, err := parseBatchArgs(c.args)
|
||||
if c.ok {
|
||||
if err != nil {
|
||||
t.Errorf("%v: unexpected error %v", c.args, err)
|
||||
continue
|
||||
}
|
||||
if verb != c.verb || !slices.Equal(rest, c.rest) {
|
||||
t.Errorf("%v: got verb=%q rest=%q", c.args, verb, rest)
|
||||
}
|
||||
} else if err == nil {
|
||||
t.Errorf("%v: want an error", c.args)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
package config
|
||||
|
||||
// P6 task 08: the [batch] block and ParseSize.
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBatchBlock(t *testing.T) {
|
||||
p := filepath.Join(t.TempDir(), "sift.toml")
|
||||
body := "[batch]\ndisabled = false\ntargets_per_job = 16\njobs_per_caller = 2\njobs_total = 4\nworkers = 2\nttl = \"12h\"\nmax_bytes = \"32MiB\"\n"
|
||||
if err := os.WriteFile(p, []byte(body), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
f, err := Load(p)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
b := f.Batch
|
||||
if b.Disabled || b.TargetsPerJob != 16 || b.JobsPerCaller != 2 || b.JobsTotal != 4 || b.Workers != 2 || b.TTL != "12h" || b.MaxBytes != "32MiB" {
|
||||
t.Fatalf("batch block %+v", b)
|
||||
}
|
||||
// No block: zero values, and batch is on (the runner applies its own defaults).
|
||||
empty := filepath.Join(t.TempDir(), "empty.toml")
|
||||
if err := os.WriteFile(empty, []byte("listen = \":8443\"\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if f, err := Load(empty); err != nil || f.Batch != (Batch{}) {
|
||||
t.Fatalf("no [batch] block: %+v %v", f.Batch, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseSize(t *testing.T) {
|
||||
const def = 7
|
||||
for _, c := range []struct {
|
||||
in string
|
||||
want int64
|
||||
}{
|
||||
{"64MiB", 64 << 20},
|
||||
{"512KiB", 512 << 10},
|
||||
{"1GiB", 1 << 30},
|
||||
{"1000", 1000},
|
||||
{" 2MiB ", 2 << 20},
|
||||
{"", def},
|
||||
{"junk", def},
|
||||
{"-1MiB", def},
|
||||
{"0", def},
|
||||
{"12MB", def}, // decimal units are not accepted: say MiB
|
||||
} {
|
||||
if got := ParseSize(c.in, def); got != c.want {
|
||||
t.Errorf("ParseSize(%q) = %d, want %d", c.in, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,254 @@
|
||||
package dcache
|
||||
|
||||
import (
|
||||
"io"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func newBatchDB(t *testing.T) *DB {
|
||||
t.Helper()
|
||||
db, err := Open("", 1<<20, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
return db
|
||||
}
|
||||
|
||||
func strp(s string) *string { return &s }
|
||||
|
||||
func TestBatchJobRoundTrip(t *testing.T) {
|
||||
db := newBatchDB(t)
|
||||
j := BatchJob{
|
||||
ID: "j1", Caller: "alice", CallerClass: ClassFleet, IDemKey: "k1",
|
||||
Status: "queued", TargetsTotal: 3, CreatedAt: 100, Lease: 100,
|
||||
}
|
||||
if err := db.CreateBatchJob(j); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, ok := db.BatchJob("j1")
|
||||
if !ok {
|
||||
t.Fatal("missing job")
|
||||
}
|
||||
if got.ID != j.ID || got.Caller != "alice" || got.CallerClass != ClassFleet ||
|
||||
got.Status != "queued" || got.TargetsTotal != 3 || got.Lease != 100 || got.CreatedAt != 100 {
|
||||
t.Fatalf("round trip mismatch: %+v", got)
|
||||
}
|
||||
if _, ok := db.BatchJob("nope"); ok {
|
||||
t.Fatal("unknown id must be ok=false")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchIDemLookup(t *testing.T) {
|
||||
db := newBatchDB(t)
|
||||
if err := db.CreateBatchJob(BatchJob{ID: "j1", Caller: "a", CallerClass: ClassFleet, IDemKey: "k", Status: "running", TargetsTotal: 1, CreatedAt: 100, Lease: 1000}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if id, ok := db.BatchJobByIDem("a", "k", 2000, time.Hour); !ok || id != "j1" {
|
||||
t.Fatalf("live idem: %q %v", id, ok)
|
||||
}
|
||||
if _, ok := db.BatchJobByIDem("a", "k", 5000, time.Hour); ok {
|
||||
t.Fatal("lease-expired job must not match")
|
||||
}
|
||||
if _, ok := db.BatchJobByIDem("b", "k", 2000, time.Hour); ok {
|
||||
t.Fatal("other caller must not match")
|
||||
}
|
||||
// Only a job that is live or has results is reused. A failed (e.g. killed by a restart)
|
||||
// or cancelled job must not be handed back to a resubmit for the rest of its lease.
|
||||
for _, c := range []struct {
|
||||
status string
|
||||
reuse bool
|
||||
}{{"queued", true}, {"done", true}, {"partial", true}, {"failed", false}, {"cancelled", false}} {
|
||||
key := "k-" + c.status
|
||||
if err := db.CreateBatchJob(BatchJob{ID: "j-" + c.status, Caller: "a", CallerClass: ClassFleet, IDemKey: key, Status: c.status, TargetsTotal: 1, CreatedAt: 100, Lease: 1000}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, ok := db.BatchJobByIDem("a", key, 2000, time.Hour); ok != c.reuse {
|
||||
t.Errorf("status %s: reused=%v, want %v", c.status, ok, c.reuse)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchLiveCounts(t *testing.T) {
|
||||
db := newBatchDB(t)
|
||||
mk := func(id, caller, status string) {
|
||||
if err := db.CreateBatchJob(BatchJob{ID: id, Caller: caller, CallerClass: ClassFleet, IDemKey: id, Status: status, TargetsTotal: 1, CreatedAt: 1, Lease: 1}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
mk("j1", "a", "queued")
|
||||
mk("j2", "a", "running")
|
||||
mk("j3", "a", "done")
|
||||
mk("j4", "b", "running")
|
||||
if n := db.LiveBatchJobs("a"); n != 2 {
|
||||
t.Fatalf("caller a live=%d want 2", n)
|
||||
}
|
||||
if n := db.LiveBatchJobsTotal(); n != 3 {
|
||||
t.Fatalf("total live=%d want 3", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchTargetsOrderedAndUpdate(t *testing.T) {
|
||||
db := newBatchDB(t)
|
||||
must := func(err error) {
|
||||
t.Helper()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
must(db.CreateBatchJob(BatchJob{ID: "j1", Caller: "a", CallerClass: ClassFleet, IDemKey: "k", Status: "running", TargetsTotal: 3, CreatedAt: 100, Lease: 100}))
|
||||
must(db.PutBatchTargets("j1", []BatchTarget{
|
||||
{JobID: "j1", Idx: 0, URL: "https://a.example/", Options: `{"mode":"clean"}`, Status: "pending"},
|
||||
{JobID: "j1", Idx: 1, URL: "https://b.example/", Status: "pending"},
|
||||
{JobID: "j1", Idx: 2, URL: "https://c.example/", Status: "pending"},
|
||||
}))
|
||||
got, ok := db.BatchTargets("j1")
|
||||
if !ok || len(got) != 3 {
|
||||
t.Fatalf("targets len=%d ok=%v", len(got), ok)
|
||||
}
|
||||
for i := range got {
|
||||
if got[i].Idx != i {
|
||||
t.Fatalf("order broken: got[%d].Idx=%d", i, got[i].Idx)
|
||||
}
|
||||
}
|
||||
|
||||
// A done write bumps the lease to now and counts the target once.
|
||||
res := strp(`{"content":"ok"}`)
|
||||
must(db.UpdateBatchTarget("j1", 0, 200, "done", res, nil, nil))
|
||||
j, _ := db.BatchJob("j1")
|
||||
if j.Lease != 200 || j.TargetsDone != 1 {
|
||||
t.Fatalf("after first write lease=%d done=%d", j.Lease, j.TargetsDone)
|
||||
}
|
||||
// Re-writing the same target (still a done transition) must not double-count.
|
||||
must(db.UpdateBatchTarget("j1", 0, 201, "done", res, nil, nil))
|
||||
j, _ = db.BatchJob("j1")
|
||||
if j.TargetsDone != 1 {
|
||||
t.Fatalf("double-counted: done=%d", j.TargetsDone)
|
||||
}
|
||||
// A failed write counts too.
|
||||
must(db.UpdateBatchTarget("j1", 1, 202, "failed", nil, strp("upstream 500"), strp("browser")))
|
||||
j, _ = db.BatchJob("j1")
|
||||
if j.TargetsDone != 2 || j.Lease != 202 {
|
||||
t.Fatalf("after failed write lease=%d done=%d", j.Lease, j.TargetsDone)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchTouchAndStatus(t *testing.T) {
|
||||
db := newBatchDB(t)
|
||||
if err := db.CreateBatchJob(BatchJob{ID: "j1", Caller: "a", CallerClass: ClassFleet, IDemKey: "k", Status: "running", TargetsTotal: 1, CreatedAt: 100, Lease: 100}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.TouchBatchJob("j1", 500); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if j, _ := db.BatchJob("j1"); j.Lease != 500 {
|
||||
t.Fatalf("touch lease=%d want 500", j.Lease)
|
||||
}
|
||||
if err := db.SetBatchJobStatus("j1", "cancelled", "", 600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if j, _ := db.BatchJob("j1"); j.Status != "cancelled" || j.Lease != 600 {
|
||||
t.Fatalf("status=%q lease=%d", j.Status, j.Lease)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSweepBatch(t *testing.T) {
|
||||
db := newBatchDB(t)
|
||||
mk := func(id string, lease int64) {
|
||||
if err := db.CreateBatchJob(BatchJob{ID: id, Caller: "a", CallerClass: ClassFleet, IDemKey: id, Status: "done", TargetsTotal: 1, CreatedAt: 1, Lease: lease}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.PutBatchTargets(id, []BatchTarget{{JobID: id, Idx: 0, URL: "https://x.example/", Status: "done", Result: strp("x")}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
mk("old", 100)
|
||||
mk("new", 1000)
|
||||
n := db.SweepBatch(300, 100*time.Second) // old: 100+100=200<=300 expired; new: 1100>300 kept
|
||||
if n != 1 {
|
||||
t.Fatalf("swept %d want 1", n)
|
||||
}
|
||||
if _, ok := db.BatchJob("old"); ok {
|
||||
t.Fatal("old job must be gone")
|
||||
}
|
||||
if _, ok := db.BatchJob("new"); !ok {
|
||||
t.Fatal("new job must remain")
|
||||
}
|
||||
if _, ok := db.BatchTargets("old"); ok {
|
||||
t.Fatal("old targets must be gone")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEvictBatchOldest(t *testing.T) {
|
||||
db := newBatchDB(t)
|
||||
done := func(id string, lease int64, size int) {
|
||||
if err := db.CreateBatchJob(BatchJob{ID: id, Caller: "a", CallerClass: ClassFleet, IDemKey: id, Status: "done", TargetsTotal: 1, CreatedAt: 1, Lease: lease}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.PutBatchTargets(id, []BatchTarget{{JobID: id, Idx: 0, URL: "https://x.example/", Status: "pending"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.UpdateBatchTarget(id, 0, lease, "done", strp(strings.Repeat("x", size)), nil, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
done("old1", 100, 100)
|
||||
done("old2", 200, 200)
|
||||
done("new", 300, 300)
|
||||
|
||||
// A running job with one completed target must never be byte-evicted.
|
||||
if err := db.CreateBatchJob(BatchJob{ID: "run", Caller: "a", CallerClass: ClassFleet, IDemKey: "run", Status: "running", TargetsTotal: 2, TargetsDone: 1, CreatedAt: 1, Lease: 400}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.PutBatchTargets("run", []BatchTarget{
|
||||
{JobID: "run", Idx: 0, URL: "https://x.example/", Status: "pending"},
|
||||
{JobID: "run", Idx: 1, URL: "https://y.example/", Status: "pending"},
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.UpdateBatchTarget("run", 0, 400, "done", strp(strings.Repeat("x", 500)), nil, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// total nbytes = 100+200+300+500 = 1100. Evict oldest done until <= 250:
|
||||
// old1(100), old2(200), new(300) all evicted (oldest first); run stays (live).
|
||||
db.EvictBatchOldest(250)
|
||||
for _, id := range []string{"old1", "old2", "new"} {
|
||||
if _, ok := db.BatchJob(id); ok {
|
||||
t.Fatalf("%s should be byte-evicted", id)
|
||||
}
|
||||
}
|
||||
if _, ok := db.BatchJob("run"); !ok {
|
||||
t.Fatal("running job must not be byte-evicted")
|
||||
}
|
||||
if n := db.BatchNbytes(); n != 500 {
|
||||
t.Fatalf("BatchNbytes=%d want 500", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMarkInFlightFailed(t *testing.T) {
|
||||
db := newBatchDB(t)
|
||||
mk := func(id, status string) {
|
||||
if err := db.CreateBatchJob(BatchJob{ID: id, Caller: "a", CallerClass: ClassFleet, IDemKey: id, Status: status, TargetsTotal: 1, CreatedAt: 1, Lease: 1}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
mk("q", "queued")
|
||||
mk("r", "running")
|
||||
mk("d", "done")
|
||||
if n := db.MarkInFlightFailed("restart", 999); n != 2 {
|
||||
t.Fatalf("marked %d want 2", n)
|
||||
}
|
||||
if j, _ := db.BatchJob("q"); j.Status != "failed" || j.Error != "restart" || j.Lease != 999 {
|
||||
t.Fatalf("q %+v", j)
|
||||
}
|
||||
if j, _ := db.BatchJob("r"); j.Status != "failed" || j.Error != "restart" {
|
||||
t.Fatalf("r %+v", j)
|
||||
}
|
||||
if j, _ := db.BatchJob("d"); j.Status != "done" {
|
||||
t.Fatalf("done job changed: %+v", j)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
package metrics
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type batchSrc struct{}
|
||||
|
||||
func (batchSrc) Snapshot() Snapshot { return Snapshot{} }
|
||||
|
||||
func TestBatchSeries(t *testing.T) {
|
||||
s := New(batchSrc{})
|
||||
s.ObserveBatchJob("created")
|
||||
s.ObserveBatchJob("done")
|
||||
s.ObserveBatchJob("done")
|
||||
s.ObserveBatchJob("reused")
|
||||
s.ObserveBatchJob("some raw string") // folded to "other", never a label
|
||||
s.ObserveBatchTarget("ok")
|
||||
s.ObserveBatchTarget("ok")
|
||||
s.ObserveBatchTarget("blocked")
|
||||
s.ObserveBatchTarget("upstream_5xx")
|
||||
s.ObserveBatchTarget("leak me") // folded to "other"
|
||||
s.SetBatchQueueDepth(3)
|
||||
s.SetBatchTargetsInflight(1)
|
||||
s.ObserveBatchJobSeconds(4 * time.Second)
|
||||
s.ObserveBatchTargetSeconds(1500 * time.Millisecond)
|
||||
s.ObserveBatchResultBytes(4096)
|
||||
|
||||
req := httptest.NewRequest("GET", "/metrics", nil)
|
||||
req.RemoteAddr, req.Host = "127.0.0.1:9", "127.0.0.1:8890"
|
||||
rec := httptest.NewRecorder()
|
||||
s.Handler().ServeHTTP(rec, req)
|
||||
body, _ := io.ReadAll(rec.Body)
|
||||
out := string(body)
|
||||
for _, want := range []string{
|
||||
`sift_batch_jobs_total{outcome="created"} 1`,
|
||||
`sift_batch_jobs_total{outcome="done"} 2`,
|
||||
`sift_batch_jobs_total{outcome="reused"} 1`,
|
||||
`sift_batch_jobs_total{outcome="other"} 1`,
|
||||
`sift_batch_targets_total{outcome="ok"} 2`,
|
||||
`sift_batch_targets_total{outcome="blocked"} 1`,
|
||||
`sift_batch_targets_total{outcome="upstream_5xx"} 1`,
|
||||
`sift_batch_targets_total{outcome="other"} 1`,
|
||||
`sift_batch_queue_depth 3`,
|
||||
`sift_batch_targets_inflight 1`,
|
||||
`sift_batch_job_seconds_count 1`,
|
||||
`sift_batch_target_seconds_count 1`,
|
||||
`sift_batch_result_bytes_count 1`,
|
||||
} {
|
||||
if !strings.Contains(out, want) {
|
||||
t.Errorf("metrics lack %q", want)
|
||||
}
|
||||
}
|
||||
if strings.Contains(out, "raw") || strings.Contains(out, "leak") {
|
||||
t.Error("a raw outcome string leaked into a label")
|
||||
}
|
||||
}
|
||||
+75
-218
@@ -6,7 +6,6 @@ import (
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
@@ -16,9 +15,9 @@ import (
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/auth"
|
||||
"git.wntrmute.dev/kyle/sift/internal/batch"
|
||||
"git.wntrmute.dev/kyle/sift/internal/clean"
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
"git.wntrmute.dev/kyle/sift/internal/digest"
|
||||
@@ -29,7 +28,6 @@ import (
|
||||
"git.wntrmute.dev/kyle/sift/internal/pdfdoc"
|
||||
"git.wntrmute.dev/kyle/sift/internal/pdfpool"
|
||||
"git.wntrmute.dev/kyle/sift/internal/pdftext"
|
||||
"git.wntrmute.dev/kyle/sift/internal/recover"
|
||||
"git.wntrmute.dev/kyle/sift/internal/search"
|
||||
"git.wntrmute.dev/kyle/sift/internal/selfip"
|
||||
)
|
||||
@@ -38,6 +36,14 @@ const (
|
||||
maxReqBody = 64 << 10
|
||||
writeTimeout = 240 * time.Second // P5a: a PDF may take 120 s to download + 20 s queued + 60 s to extract
|
||||
requestDeadline = 210 * time.Second // P5a: the PDF path's worst case (see writeTimeout); other paths keep their own shorter timeouts
|
||||
|
||||
// batchMaxBody is the largest body POST /v1/batch accepts: 32 targets, each with a
|
||||
// large question, exceed the 64 KiB cap of the other endpoints.
|
||||
batchMaxBody = 256 << 10
|
||||
|
||||
// DefaultPDFMaxUpload is the largest body POST /v1/extract accepts: an uploaded
|
||||
// PDF, addressed by its hash and extracted once per content.
|
||||
DefaultPDFMaxUpload = 100 << 20
|
||||
)
|
||||
|
||||
type Server struct {
|
||||
@@ -55,6 +61,7 @@ type Server struct {
|
||||
PDFPool *pdfpool.Pool
|
||||
PDFTTL time.Duration
|
||||
PDFMaxQuestionChunks int
|
||||
PDFMaxUpload int64 // 0 = DefaultPDFMaxUpload; the /v1/extract upload cap
|
||||
Search *search.Engine
|
||||
Ctx *fleetctx.Client
|
||||
// PermitRestricted is context.permit_restricted: those callers never
|
||||
@@ -69,6 +76,14 @@ type Server struct {
|
||||
reqN, blockN, scrapeN atomic.Uint64
|
||||
|
||||
pdfPoolOnce sync.Once
|
||||
|
||||
// Batch is the batch runner; nil means the /v1/batch endpoints (task 07) answer 503
|
||||
// "batch disabled".
|
||||
Batch *batch.Runner
|
||||
|
||||
// uploads bounds how many /v1/extract bodies (up to PDFMaxUpload each) are read into
|
||||
// memory at once, ahead of the PDF pool's own limits.
|
||||
uploads chan struct{}
|
||||
}
|
||||
|
||||
func New(version string, g *guard.Guard, f *fetch.Fetcher, keys *auth.Store, log *slog.Logger) *Server {
|
||||
@@ -85,6 +100,7 @@ func New(version string, g *guard.Guard, f *fetch.Fetcher, keys *auth.Store, log
|
||||
OutCap: 200_000,
|
||||
global: make(chan struct{}, 8),
|
||||
per: map[string]chan struct{}{},
|
||||
uploads: make(chan struct{}, maxUploads),
|
||||
}
|
||||
s.Prom = metrics.New(s)
|
||||
s.Mux.HandleFunc("GET /healthz", s.healthz)
|
||||
@@ -93,8 +109,13 @@ func New(version string, g *guard.Guard, f *fetch.Fetcher, keys *auth.Store, log
|
||||
s.Mux.HandleFunc("GET /metrics.json", s.metrics)
|
||||
s.Mux.HandleFunc("POST /v2/scrape", s.scrape)
|
||||
s.Mux.HandleFunc("POST /v1/fetch", s.v1fetch)
|
||||
s.Mux.HandleFunc("POST /v1/extract", s.v1extract)
|
||||
s.Mux.HandleFunc("POST /v1/search", s.v1search)
|
||||
s.Mux.HandleFunc("POST /v2/search", s.v2search)
|
||||
s.Mux.HandleFunc("POST /v1/batch", s.batchSubmit)
|
||||
s.Mux.HandleFunc("GET /v1/batch/{id}", s.batchStatus)
|
||||
s.Mux.HandleFunc("GET /v1/batch/{id}/{idx}", s.batchResult)
|
||||
s.Mux.HandleFunc("DELETE /v1/batch/{id}", s.batchCancel)
|
||||
return s
|
||||
}
|
||||
|
||||
@@ -111,9 +132,16 @@ func (s *Server) HTTPServer(addr string) *http.Server {
|
||||
}
|
||||
|
||||
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
r.Body = http.MaxBytesReader(w, r.Body, maxReqBody)
|
||||
var limit int64 = maxReqBody
|
||||
switch r.URL.Path {
|
||||
case "/v1/fetch", "/v2/scrape", "/v1/search", "/v2/search":
|
||||
case "/v1/extract":
|
||||
limit = int64(s.pdfMaxUpload())
|
||||
case "/v1/batch":
|
||||
limit = batchMaxBody
|
||||
}
|
||||
r.Body = http.MaxBytesReader(w, r.Body, limit)
|
||||
switch r.URL.Path {
|
||||
case "/v1/fetch", "/v2/scrape", "/v1/search", "/v2/search", "/v1/extract":
|
||||
ctx, cancel := context.WithTimeout(r.Context(), requestDeadline)
|
||||
defer cancel()
|
||||
r = r.WithContext(ctx)
|
||||
@@ -127,9 +155,29 @@ func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
s.Prom.Observe(r.URL.Path, time.Since(start), code)
|
||||
return
|
||||
}
|
||||
if strings.HasPrefix(r.URL.Path, "/v1/batch") {
|
||||
cw := &codeWriter{ResponseWriter: w}
|
||||
start := time.Now()
|
||||
s.Mux.ServeHTTP(cw, r)
|
||||
code := cw.code
|
||||
if code == 0 {
|
||||
code = http.StatusOK
|
||||
}
|
||||
s.Prom.Observe(batchObservePath(r.URL.Path), time.Since(start), code)
|
||||
return
|
||||
}
|
||||
s.Mux.ServeHTTP(w, r)
|
||||
}
|
||||
|
||||
// batchObservePath collapses a batch path to a fixed label so a job id never becomes a
|
||||
// metric label value: the submit path stays, every {id}/{idx} variant becomes {id}.
|
||||
func batchObservePath(path string) string {
|
||||
if path == "/v1/batch" {
|
||||
return "/v1/batch"
|
||||
}
|
||||
return "/v1/batch/{id}"
|
||||
}
|
||||
|
||||
type codeWriter struct {
|
||||
http.ResponseWriter
|
||||
code int
|
||||
@@ -374,7 +422,7 @@ func (s *Server) scrape(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
}
|
||||
s.scrapeN.Add(1)
|
||||
res, cl, doc, ferr := s.doFetch(r, c, req.URL, "", onlyMain)
|
||||
res, cl, doc, ferr := s.doFetch(r.Context(), c, req.URL, "", onlyMain)
|
||||
if ferr != nil {
|
||||
s.writeFetchErr(w, c, req.URL, ferr)
|
||||
return
|
||||
@@ -434,220 +482,29 @@ func (s *Server) v1fetch(w http.ResponseWriter, r *http.Request) {
|
||||
writeErr(w, 400, "invalid url", "")
|
||||
return
|
||||
}
|
||||
q := req.Question
|
||||
if len([]rune(q)) > 500 {
|
||||
rns := []rune(q)
|
||||
q = string(rns[:500])
|
||||
t := batch.Target{
|
||||
URL: req.URL,
|
||||
Options: batch.Options{
|
||||
Mode: req.Mode,
|
||||
Question: req.Question,
|
||||
Pages: req.Pages,
|
||||
Grep: req.Grep,
|
||||
JQ: req.JQ,
|
||||
Context: req.Context,
|
||||
Recover: req.Recover,
|
||||
},
|
||||
}
|
||||
mode := req.Mode
|
||||
if mode == "" {
|
||||
if q != "" {
|
||||
mode = "digest"
|
||||
} else {
|
||||
mode = "clean"
|
||||
}
|
||||
}
|
||||
if mode != "clean" && mode != "raw" && mode != "digest" {
|
||||
writeErr(w, 400, "invalid url", "")
|
||||
return
|
||||
}
|
||||
wantCtx := q != ""
|
||||
if req.Context != nil {
|
||||
wantCtx = *req.Context
|
||||
}
|
||||
|
||||
t0 := time.Now()
|
||||
onlyMain := true
|
||||
res, cl, doc, ferr := s.doFetch(r, c, req.URL, req.JQ, onlyMain)
|
||||
var snap *recover.Result
|
||||
if ferr != nil && req.Recover && recoverWorthy(ferr) {
|
||||
sr, rerr := recover.Try(r.Context(), s.Fetch, req.URL)
|
||||
if rerr == nil && sr != nil && sr.Res != nil {
|
||||
res, cl, doc, ferr = s.cleanFetched(sr.Res, req.JQ, onlyMain)
|
||||
snap = sr
|
||||
}
|
||||
}
|
||||
if ferr != nil {
|
||||
s.writeFetchErr(w, c, req.URL, ferr)
|
||||
return
|
||||
}
|
||||
fetchMS := time.Since(t0).Milliseconds()
|
||||
|
||||
if doc == nil && (req.Pages != "" || req.Grep != "") {
|
||||
s.writeFetchErr(w, c, req.URL, &fetch.Error{Status: 400, Msg: "pages and grep apply to PDFs only"})
|
||||
return
|
||||
}
|
||||
|
||||
var docLen, pdfPages int
|
||||
var content string
|
||||
var pdfMap map[string]any
|
||||
if doc != nil {
|
||||
pdfPages = doc.Pages
|
||||
if mode == "raw" {
|
||||
writeErr(w, 400, "raw mode does not apply to PDFs", "")
|
||||
return
|
||||
}
|
||||
rp, err := s.renderPDF(doc, pdfdoc.Options{Pages: req.Pages, Grep: req.Grep, OutCap: s.OutCap})
|
||||
res, err := s.fetchTarget(r.Context(), c, t)
|
||||
if err != nil {
|
||||
if errors.Is(err, pdfdoc.ErrBadPages) || errors.Is(err, pdfdoc.ErrBadGrep) {
|
||||
s.writeFetchErr(w, c, req.URL, &fetch.Error{Status: 400, Msg: err.Error()})
|
||||
return
|
||||
}
|
||||
s.writeFetchErr(w, c, req.URL, err)
|
||||
return
|
||||
}
|
||||
full := pdfdoc.Full(doc)
|
||||
docLen = utf8.RuneCountInString(full)
|
||||
if mode == "digest" {
|
||||
content = full
|
||||
} else {
|
||||
content = rp.Content
|
||||
env := make(map[string]any, len(res.Meta)+1)
|
||||
for k, v := range res.Meta {
|
||||
env[k] = v
|
||||
}
|
||||
pdfMap = map[string]any{
|
||||
"pages": doc.Pages,
|
||||
"title": doc.Title,
|
||||
"outline": doc.Outline,
|
||||
"extracted": len(doc.Text),
|
||||
"first": rp.First,
|
||||
"last": rp.Last,
|
||||
"map": rp.Map,
|
||||
"grep_hits": rp.GrepHits,
|
||||
"note": rp.Note,
|
||||
}
|
||||
} else {
|
||||
content = cl.Markdown
|
||||
if cl.Kind == "json" {
|
||||
content = cl.HTML
|
||||
}
|
||||
if mode == "raw" {
|
||||
content = string(res.Body)
|
||||
if len(content) > s.OutCap {
|
||||
content = clean.TruncateUTF8(content, s.OutCap)
|
||||
cl.Trunc = true
|
||||
cl.Warning = "truncated to output cap"
|
||||
}
|
||||
}
|
||||
}
|
||||
warn := []string{}
|
||||
if cl.Warning != "" {
|
||||
warn = append(warn, cl.Warning)
|
||||
}
|
||||
|
||||
var digMeta map[string]any
|
||||
var digestMS int64
|
||||
if mode == "digest" {
|
||||
if doc != nil {
|
||||
n := len(digest.Split(pdfdoc.Full(doc), digest.ChunkTokens))
|
||||
limit := s.PDFMaxQuestionChunks
|
||||
if limit == 0 {
|
||||
limit = DefaultPDFMaxQuestionChunks
|
||||
}
|
||||
if n > limit {
|
||||
s.writeFetchErr(w, c, req.URL, &fetch.Error{Status: 422, Msg: fmt.Sprintf("document too long for question (%d chunks); use grep or pages", n)})
|
||||
return
|
||||
}
|
||||
}
|
||||
if s.Digest == nil {
|
||||
warn = append(warn, "digest fell back to clean: unset")
|
||||
mode = "clean"
|
||||
content, warn = capDigestFallback(content, warn)
|
||||
digMeta = map[string]any{"fallback": "unset"}
|
||||
} else {
|
||||
dr := s.Digest.Run(r.Context(), content, q, res.FinalURL, s.callerClass(c))
|
||||
digestMS = dr.DigestMS
|
||||
if !dr.Cached {
|
||||
s.Prom.ObserveDigest(time.Duration(dr.DigestMS) * time.Millisecond)
|
||||
}
|
||||
if dr.Fallback != "" {
|
||||
warn = append(warn, "digest fell back to clean: "+dr.Fallback)
|
||||
mode = "clean"
|
||||
content, warn = capDigestFallback(content, warn)
|
||||
digMeta = map[string]any{"fallback": dr.Fallback, "model": dr.Model}
|
||||
} else {
|
||||
content = dr.Text
|
||||
digMeta = map[string]any{
|
||||
"model": dr.Model,
|
||||
"chunks_used": dr.ChunksUsed,
|
||||
"chunks_total": dr.ChunksTotal,
|
||||
"fallback": false,
|
||||
"cached": dr.Cached,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var ctxItems []fleetctx.Item
|
||||
var ctxMS int64
|
||||
if wantCtx && s.Ctx != nil {
|
||||
ct0 := time.Now()
|
||||
query := fleetctx.Query(q, cl.Title, headings(cl.Markdown))
|
||||
items, cw := s.Ctx.Search(r.Context(), c.Workspace, c.Name, query)
|
||||
ctxMS = time.Since(ct0).Milliseconds()
|
||||
if cw != "" {
|
||||
warn = append(warn, cw)
|
||||
}
|
||||
ctxItems = items
|
||||
}
|
||||
|
||||
qLen := len(q)
|
||||
if qLen > 80 {
|
||||
qLen = 80
|
||||
}
|
||||
fb, _ := "", false
|
||||
if digMeta != nil {
|
||||
if s, ok := digMeta["fallback"].(string); ok {
|
||||
fb = s
|
||||
}
|
||||
}
|
||||
s.Log.Info("fetch", "caller", c.Name, "host", s.logHost(c, res.FinalURL), "mode", mode,
|
||||
"q_len", qLen, "raw_tokens", len(res.Body)/4, "returned_tokens", len(content)/4,
|
||||
"pdf_pages", pdfPages, "digest_ms", digestMS, "fallback", fb)
|
||||
|
||||
src := map[string]any{
|
||||
"url": res.URL,
|
||||
"final_url": res.FinalURL,
|
||||
"status": res.Status,
|
||||
"content_type": res.ContentType,
|
||||
"fetched_at": res.FetchedAt.Format(time.RFC3339),
|
||||
"hash": res.SHA256(),
|
||||
"bytes": len(res.Body),
|
||||
"cached": len(res.Body) == 0,
|
||||
"kind": "live",
|
||||
}
|
||||
if snap != nil {
|
||||
src["kind"] = "snapshot"
|
||||
src["via"] = string(snap.Via)
|
||||
src["original_url"] = snap.OriginalURL
|
||||
if snap.SnapshotAt != "" {
|
||||
src["snapshot_at"] = snap.SnapshotAt
|
||||
} else {
|
||||
src["snapshot_at"] = nil
|
||||
warn = append(warn, "snapshot date unknown")
|
||||
}
|
||||
}
|
||||
out := map[string]any{
|
||||
"content": content,
|
||||
"mode": mode,
|
||||
"untrusted": true,
|
||||
"source": src,
|
||||
"warnings": warn,
|
||||
"sizes": map[string]int{
|
||||
"raw_tokens": docLen / 4,
|
||||
"clean_tokens": digest.EstTokens(cl.Markdown),
|
||||
"returned_tokens": digest.EstTokens(content),
|
||||
},
|
||||
"timings": map[string]int64{"fetch_ms": fetchMS, "digest_ms": digestMS, "context_ms": ctxMS},
|
||||
}
|
||||
if pdfMap != nil {
|
||||
out["pdf"] = pdfMap
|
||||
}
|
||||
if digMeta != nil {
|
||||
out["digest"] = digMeta
|
||||
}
|
||||
if wantCtx {
|
||||
out["context"] = ctxItems
|
||||
}
|
||||
writeJSON(w, 200, out)
|
||||
env["content"] = res.Content
|
||||
writeJSON(w, 200, env)
|
||||
}
|
||||
|
||||
func recoverWorthy(err error) bool {
|
||||
@@ -693,8 +550,8 @@ func (s *Server) cleanFetched(res *fetch.Resource, jq string, onlyMain bool) (*f
|
||||
return res, cl, nil, nil
|
||||
}
|
||||
|
||||
func (s *Server) doFetch(r *http.Request, c *auth.Caller, rawURL, jq string, onlyMain bool) (*fetch.Resource, *clean.Result, *pdftext.Doc, error) {
|
||||
res, hitDoc, ok, err := s.pdfFromCache(r.Context(), c, rawURL)
|
||||
func (s *Server) doFetch(ctx context.Context, c *auth.Caller, rawURL, jq string, onlyMain bool) (*fetch.Resource, *clean.Result, *pdftext.Doc, error) {
|
||||
res, hitDoc, ok, err := s.pdfFromCache(ctx, c, rawURL)
|
||||
if err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
@@ -707,7 +564,7 @@ func (s *Server) doFetch(r *http.Request, c *auth.Caller, rawURL, jq string, onl
|
||||
}
|
||||
defer s.release(c.Name, rawURL)
|
||||
|
||||
res, err = s.Fetch.Fetch(r.Context(), rawURL)
|
||||
res, err = s.Fetch.Fetch(ctx, rawURL)
|
||||
if err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
@@ -739,7 +596,7 @@ func (s *Server) doFetch(r *http.Request, c *auth.Caller, rawURL, jq string, onl
|
||||
}
|
||||
var doc *pdftext.Doc
|
||||
if kind == "pdf" {
|
||||
doc, err = s.extractPDF(r.Context(), c, res)
|
||||
doc, err = s.extractPDF(ctx, c, res)
|
||||
if err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
|
||||
@@ -0,0 +1,372 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/auth"
|
||||
"git.wntrmute.dev/kyle/sift/internal/batch"
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
"git.wntrmute.dev/kyle/sift/internal/fetch"
|
||||
)
|
||||
|
||||
// batchCtxKey marks a context as a batch fetch so extractPDF charges the PDF pool to the
|
||||
// caller's "/batch" slot instead of its interactive one.
|
||||
type batchCtxKey struct{}
|
||||
|
||||
func isBatchFetch(ctx context.Context) bool {
|
||||
ok, _ := ctx.Value(batchCtxKey{}).(bool)
|
||||
return ok
|
||||
}
|
||||
|
||||
// NewBatchRunner builds the runner that serve (task 07) and the tests use. Its FetchFunc is
|
||||
// s.batchFetch and its Hooks feed s.Prom; any cfg.Hooks from the caller is overwritten.
|
||||
func (s *Server) NewBatchRunner(db *dcache.DB, cfg batch.Config) *batch.Runner {
|
||||
cfg.Hooks = batch.Hooks{
|
||||
Job: func(outcome string) {
|
||||
s.Prom.ObserveBatchJob(outcome)
|
||||
},
|
||||
Target: func(err error, d time.Duration, n int) {
|
||||
s.Prom.ObserveBatchTarget(batchOutcome(err))
|
||||
s.Prom.ObserveBatchTargetSeconds(d)
|
||||
if err == nil {
|
||||
s.Prom.ObserveBatchResultBytes(n)
|
||||
}
|
||||
},
|
||||
JobSeconds: func(d time.Duration) {
|
||||
s.Prom.ObserveBatchJobSeconds(d)
|
||||
},
|
||||
Depth: func(live int) {
|
||||
s.Prom.SetBatchQueueDepth(live)
|
||||
},
|
||||
Inflight: func(n int) {
|
||||
s.Prom.SetBatchTargetsInflight(n)
|
||||
},
|
||||
}
|
||||
return batch.New(db, s.batchFetch, cfg, s.Log)
|
||||
}
|
||||
|
||||
// batchFetch is the runner's FetchFunc: one target through s.fetchTarget under the request
|
||||
// deadline, tagged as a batch fetch. Its errors map to what the runner and the metrics need.
|
||||
func (s *Server) batchFetch(ctx context.Context, c *auth.Caller, t batch.Target) (batch.Result, error) {
|
||||
ctx, cancel := context.WithTimeout(ctx, requestDeadline)
|
||||
defer cancel()
|
||||
ctx = context.WithValue(ctx, batchCtxKey{}, true)
|
||||
res, err := s.fetchTarget(ctx, c, t)
|
||||
if err != nil {
|
||||
var fe *fetch.Error
|
||||
if errors.As(err, &fe) {
|
||||
switch fe.Status {
|
||||
case 429, 503:
|
||||
return batch.Result{}, fmt.Errorf("%s: %w", fe.Msg, batch.ErrRetryLater)
|
||||
}
|
||||
// Copy, never mutate: the PDF pool's single-flight may share this error with an
|
||||
// interactive request.
|
||||
status, msg := fe.Status, fe.Msg
|
||||
if status == 502 {
|
||||
status, msg = 500, "internal error"
|
||||
}
|
||||
if status == 403 {
|
||||
s.blockN.Add(1)
|
||||
}
|
||||
return batch.Result{}, &batch.TargetError{Status: status, Msg: msg, Hint: fe.Hint}
|
||||
}
|
||||
s.Log.Error("internal", "caller", c.Name, "host", "-", "err", s.logErr(c, err))
|
||||
return batch.Result{}, &batch.TargetError{Status: 500, Msg: "internal error"}
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// batchOutcome classifies a finished target's error into the sift_batch_targets_total label:
|
||||
// a nil error is ok, a retry-later the runner gave up on is busy, a TargetError by its status,
|
||||
// and anything else upstream_5xx so a raw string never becomes a label.
|
||||
func batchOutcome(err error) string {
|
||||
switch {
|
||||
case err == nil:
|
||||
return "ok"
|
||||
case errors.Is(err, batch.ErrRetryLater):
|
||||
return "busy"
|
||||
}
|
||||
var te *batch.TargetError
|
||||
if errors.As(err, &te) {
|
||||
switch te.Status {
|
||||
case 403:
|
||||
return "blocked"
|
||||
case 429:
|
||||
return "busy"
|
||||
case 504:
|
||||
return "timeout"
|
||||
case 415:
|
||||
return "unsupported"
|
||||
case 422:
|
||||
if strings.HasPrefix(te.Msg, "document too long") {
|
||||
return "too_long"
|
||||
}
|
||||
return "upstream_4xx"
|
||||
}
|
||||
if te.Status >= 400 && te.Status < 500 {
|
||||
return "upstream_4xx"
|
||||
}
|
||||
}
|
||||
return "upstream_5xx"
|
||||
}
|
||||
|
||||
// The four /v1/batch handlers (task 07). A job that is unknown or owned by another caller is a
|
||||
// 404 for every GET and DELETE, never a 403, so a caller cannot probe for another caller's jobs.
|
||||
|
||||
// ownership answers the caller for id: a disabled batch is 503, an unknown id or one not owned by
|
||||
// c is a 404. The empty job returned on either path signals "already answered" to the caller.
|
||||
func (s *Server) ownership(w http.ResponseWriter, c *auth.Caller, id string) (batch.Job, []batch.TargetStatus) {
|
||||
if s.Batch == nil {
|
||||
writeErr(w, 503, "batch disabled", "")
|
||||
return batch.Job{}, nil
|
||||
}
|
||||
j, ts, ok := s.Batch.Status(id)
|
||||
if !ok || j.Caller != c.Name {
|
||||
writeErr(w, 404, "not found", "")
|
||||
return batch.Job{}, nil
|
||||
}
|
||||
return j, ts
|
||||
}
|
||||
|
||||
// logBatch writes the one batch line per handler: caller, job_id, targets, then either status or
|
||||
// reused. It never carries a target URL, hostname or the idem_key.
|
||||
func (s *Server) logBatch(c *auth.Caller, jobID string, targets int, extra ...any) {
|
||||
s.Log.Info("batch",
|
||||
append([]any{
|
||||
"caller", c.Name,
|
||||
"job_id", jobID,
|
||||
"targets", targets,
|
||||
}, extra...)...,
|
||||
)
|
||||
}
|
||||
|
||||
func validBatchMode(m string) bool {
|
||||
switch m {
|
||||
case "", "clean", "raw", "digest":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (s *Server) batchSubmit(w http.ResponseWriter, r *http.Request) {
|
||||
c := s.caller(w, r)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
if s.Batch == nil { // [batch] disabled, or the store did not open
|
||||
writeErr(w, 503, "batch disabled", "")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
Targets []v1Req `json:"targets"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
writeErr(w, 400, "invalid batch: bad json", "")
|
||||
return
|
||||
}
|
||||
if len(req.Targets) == 0 {
|
||||
writeErr(w, 400, "invalid batch: no targets", "")
|
||||
return
|
||||
}
|
||||
targets := make([]batch.Target, 0, len(req.Targets))
|
||||
for _, t := range req.Targets {
|
||||
if t.URL == "" {
|
||||
writeErr(w, 400, "invalid batch: empty url", "")
|
||||
return
|
||||
}
|
||||
if !validBatchMode(t.Mode) {
|
||||
writeErr(w, 400, "invalid batch: bad mode", "")
|
||||
return
|
||||
}
|
||||
targets = append(targets, batch.Target{
|
||||
URL: t.URL,
|
||||
Options: batch.Options{
|
||||
Mode: t.Mode,
|
||||
Question: t.Question,
|
||||
Pages: t.Pages,
|
||||
Grep: t.Grep,
|
||||
JQ: t.JQ,
|
||||
Context: t.Context,
|
||||
Recover: t.Recover,
|
||||
},
|
||||
})
|
||||
}
|
||||
for i, t := range targets {
|
||||
if err := s.Guard.PreCheck(t.URL); err != nil {
|
||||
s.logBatch(c, "", len(targets), "status", "blocked")
|
||||
writeErr(w, 422, fmt.Sprintf("blocked: target %d: %v", i, err), "")
|
||||
return
|
||||
}
|
||||
}
|
||||
id, reused, err := s.Batch.Submit(r.Context(), c, s.callerClass(c), targets)
|
||||
if err != nil {
|
||||
s.logBatch(c, "", len(targets), "status", "failed")
|
||||
switch {
|
||||
case errors.Is(err, batch.ErrTooManyTargets):
|
||||
writeErr(w, 422, "too many targets (max 32)", "")
|
||||
case errors.Is(err, batch.ErrCallerBusy):
|
||||
w.Header().Set("Retry-After", "30")
|
||||
writeErr(w, 429, "busy: a batch of yours is still running", "")
|
||||
case errors.Is(err, batch.ErrBusy):
|
||||
w.Header().Set("Retry-After", "60")
|
||||
writeErr(w, 429, "batch jobs busy", "")
|
||||
case errors.Is(err, batch.ErrNoTargets):
|
||||
writeErr(w, 400, "invalid batch: no targets", "")
|
||||
default:
|
||||
writeErr(w, 500, "internal error", "")
|
||||
}
|
||||
return
|
||||
}
|
||||
status := "queued"
|
||||
if reused {
|
||||
if j, _, ok := s.Batch.Status(id); ok {
|
||||
status = j.Status
|
||||
}
|
||||
}
|
||||
if reused {
|
||||
s.logBatch(c, id, len(targets), "reused", true)
|
||||
} else {
|
||||
s.logBatch(c, id, len(targets), "status", status)
|
||||
}
|
||||
writeJSON(w, 202, map[string]any{
|
||||
"success": true,
|
||||
"job_id": id,
|
||||
"status": status,
|
||||
"targets": len(targets),
|
||||
"reused": reused,
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) batchStatus(w http.ResponseWriter, r *http.Request) {
|
||||
c := s.caller(w, r)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
id := r.PathValue("id")
|
||||
j, ts := s.ownership(w, c, id)
|
||||
if j.ID == "" {
|
||||
return
|
||||
}
|
||||
s.Batch.Touch(id)
|
||||
s.logBatch(c, id, j.TargetsTotal, "status", j.Status)
|
||||
writeJSON(w, 200, statusEnvelope(j, ts, r))
|
||||
}
|
||||
|
||||
func (s *Server) batchResult(w http.ResponseWriter, r *http.Request) {
|
||||
c := s.caller(w, r)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
id := r.PathValue("id")
|
||||
j, ts := s.ownership(w, c, id)
|
||||
if j.ID == "" {
|
||||
return
|
||||
}
|
||||
idx, err := strconv.Atoi(r.PathValue("idx"))
|
||||
if err != nil || idx < 0 || idx >= j.TargetsTotal || idx >= len(ts) {
|
||||
writeErr(w, 404, "not found", "")
|
||||
return
|
||||
}
|
||||
s.Batch.Touch(id)
|
||||
t := ts[idx]
|
||||
out := map[string]any{
|
||||
"job_id": j.ID,
|
||||
"idx": t.Idx,
|
||||
"url": t.URL,
|
||||
"status": t.Status,
|
||||
}
|
||||
if t.Result != nil {
|
||||
out["result"] = resultEnvelope(t.Result)
|
||||
}
|
||||
if t.Error != "" {
|
||||
out["error"] = t.Error
|
||||
}
|
||||
if t.Hint != "" {
|
||||
out["hint"] = t.Hint
|
||||
}
|
||||
s.logBatch(c, id, j.TargetsTotal, "status", t.Status)
|
||||
writeJSON(w, 200, out)
|
||||
}
|
||||
|
||||
func (s *Server) batchCancel(w http.ResponseWriter, r *http.Request) {
|
||||
c := s.caller(w, r)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
id := r.PathValue("id")
|
||||
j, _ := s.ownership(w, c, id)
|
||||
if j.ID == "" {
|
||||
return
|
||||
}
|
||||
switch j.Status {
|
||||
case "done", "partial", "failed", "cancelled":
|
||||
s.logBatch(c, id, j.TargetsTotal, "status", "failed")
|
||||
writeErr(w, 409, "job already finished", "")
|
||||
return
|
||||
}
|
||||
if s.Batch.Cancel(id) {
|
||||
s.logBatch(c, id, j.TargetsTotal, "status", "cancelled")
|
||||
writeJSON(w, 200, map[string]any{"success": true, "status": "cancelled"})
|
||||
return
|
||||
}
|
||||
s.logBatch(c, id, j.TargetsTotal, "status", "failed")
|
||||
writeErr(w, 409, "job already finished", "")
|
||||
}
|
||||
|
||||
// statusEnvelope is the summary GET /v1/batch/{id} returns: statuses, never inline results,
|
||||
// because 32 results can be several megabytes. With ?results=1 each done target carries the
|
||||
// /v1/fetch envelope (Meta plus content).
|
||||
func statusEnvelope(j batch.Job, ts []batch.TargetStatus, r *http.Request) map[string]any {
|
||||
out := map[string]any{
|
||||
"job_id": j.ID,
|
||||
"status": j.Status,
|
||||
"targets_total": j.TargetsTotal,
|
||||
"targets_done": j.TargetsDone,
|
||||
"created_at": time.Unix(j.CreatedAt, 0).UTC().Format(time.RFC3339),
|
||||
"targets": targetSummaries(ts, r),
|
||||
}
|
||||
if j.Error != "" {
|
||||
out["error"] = j.Error
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func targetSummaries(ts []batch.TargetStatus, r *http.Request) []map[string]any {
|
||||
results := r.URL.Query().Get("results") == "1"
|
||||
out := make([]map[string]any, 0, len(ts))
|
||||
for _, t := range ts {
|
||||
m := map[string]any{
|
||||
"idx": t.Idx,
|
||||
"url": t.URL,
|
||||
"status": t.Status,
|
||||
}
|
||||
if t.Error != "" {
|
||||
m["error"] = t.Error
|
||||
}
|
||||
if t.Hint != "" {
|
||||
m["hint"] = t.Hint
|
||||
}
|
||||
if results && t.Status == "done" && t.Result != nil {
|
||||
m["result"] = resultEnvelope(t.Result)
|
||||
}
|
||||
out = append(out, m)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// resultEnvelope is the /v1/fetch shape: every Meta field plus content.
|
||||
func resultEnvelope(res *batch.Result) map[string]any {
|
||||
env := make(map[string]any, len(res.Meta)+1)
|
||||
for k, v := range res.Meta {
|
||||
env[k] = v
|
||||
}
|
||||
env["content"] = res.Content
|
||||
return env
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package api
|
||||
|
||||
// P6 owner review: with batch off (disabled, or the store did not open) every batch endpoint
|
||||
// answers 503; none of them panics on the nil runner.
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBatchDisabled(t *testing.T) {
|
||||
s, keys := newBatchServer(t, "test")
|
||||
s.Batch.Close()
|
||||
s.Batch = nil
|
||||
for _, c := range []struct{ method, path string }{
|
||||
{http.MethodPost, "/v1/batch"},
|
||||
{http.MethodGet, "/v1/batch/x"},
|
||||
{http.MethodGet, "/v1/batch/x/0"},
|
||||
{http.MethodDelete, "/v1/batch/x"},
|
||||
} {
|
||||
var body any
|
||||
if c.method == http.MethodPost {
|
||||
body = map[string]any{"targets": []any{map[string]any{"url": "http://127.0.0.1/"}}}
|
||||
}
|
||||
if rec := batchReq(t, s, c.method, c.path, keys["test"], body); rec.Code != 503 {
|
||||
t.Errorf("%s %s: %d %s, want 503", c.method, c.path, rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,280 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Every test fetches only from local httptest servers or internal.test (fake DNS): no test
|
||||
// touches the network. newBatchServer, pageServer and fakeDNS come from batchrun_test.go.
|
||||
|
||||
func batchReq(t *testing.T, s *Server, method, path, key string, body any) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
var rd io.Reader
|
||||
if body != nil {
|
||||
b, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
rd = strings.NewReader(string(b))
|
||||
}
|
||||
req := httptest.NewRequest(method, path, rd)
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
if body != nil {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
s.ServeHTTP(rec, req)
|
||||
return rec
|
||||
}
|
||||
|
||||
func batchDecode(t *testing.T, rec *httptest.ResponseRecorder) map[string]any {
|
||||
t.Helper()
|
||||
var m map[string]any
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &m); err != nil {
|
||||
t.Fatalf("decode: %v (%d %s)", err, rec.Code, rec.Body.String())
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func submit(t *testing.T, s *Server, key string, targets ...map[string]any) map[string]any {
|
||||
t.Helper()
|
||||
rec := batchReq(t, s, http.MethodPost, "/v1/batch", key, map[string]any{"targets": targets})
|
||||
if rec.Code != 202 {
|
||||
t.Fatalf("submit: %d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
return batchDecode(t, rec)
|
||||
}
|
||||
|
||||
func waitTerminal(t *testing.T, s *Server, key, id string) map[string]any {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(10 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
env := batchDecode(t, batchReq(t, s, http.MethodGet, "/v1/batch/"+id, key, nil))
|
||||
switch env["status"] {
|
||||
case "done", "partial", "failed", "cancelled":
|
||||
return env
|
||||
}
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("batch never reached a terminal status")
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestBatchSubmitAndCollect(t *testing.T) {
|
||||
ts := pageServer(t, nil)
|
||||
s, keys := newBatchServer(t, "test")
|
||||
sub := submit(t, s, keys["test"], map[string]any{"url": ts.URL + "/one"}, map[string]any{"url": ts.URL + "/two"})
|
||||
id, _ := sub["job_id"].(string)
|
||||
if sub["success"] != true || id == "" || sub["reused"] != false || sub["targets"] != float64(2) || sub["status"] != "queued" {
|
||||
t.Fatalf("submit body %v", sub)
|
||||
}
|
||||
|
||||
// The status body is a summary: statuses, no results.
|
||||
env := waitTerminal(t, s, keys["test"], id)
|
||||
if env["status"] != "done" || env["targets_total"] != float64(2) || env["targets_done"] != float64(2) {
|
||||
t.Fatalf("status %v", env)
|
||||
}
|
||||
if _, err := time.Parse(time.RFC3339, fmt.Sprint(env["created_at"])); err != nil {
|
||||
t.Fatalf("created_at %v is not RFC 3339", env["created_at"])
|
||||
}
|
||||
targets := env["targets"].([]any)
|
||||
if len(targets) != 2 {
|
||||
t.Fatalf("targets %d", len(targets))
|
||||
}
|
||||
for i, x := range targets {
|
||||
tg := x.(map[string]any)
|
||||
if tg["idx"] != float64(i) || tg["status"] != "done" {
|
||||
t.Fatalf("target %d: %v", i, tg)
|
||||
}
|
||||
if _, has := tg["result"]; has {
|
||||
t.Fatalf("summary must not carry results: target %d has one", i)
|
||||
}
|
||||
}
|
||||
|
||||
// One target's result: the /v1/fetch envelope.
|
||||
rec := batchReq(t, s, http.MethodGet, "/v1/batch/"+id+"/1", keys["test"], nil)
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("GET idx: %d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
one := batchDecode(t, rec)
|
||||
res, _ := one["result"].(map[string]any)
|
||||
if one["idx"] != float64(1) || one["status"] != "done" || res == nil {
|
||||
t.Fatalf("GET idx body %v", one)
|
||||
}
|
||||
if c, _ := res["content"].(string); !strings.Contains(c, "hello batch /two") {
|
||||
t.Fatalf("result content %q", c)
|
||||
}
|
||||
if res["untrusted"] != true || res["mode"] != "clean" {
|
||||
t.Fatalf("result is not the /v1/fetch envelope: %v", res)
|
||||
}
|
||||
if src, _ := res["source"].(map[string]any); src == nil || src["final_url"] != ts.URL+"/two" {
|
||||
t.Fatalf("result source %v", res["source"])
|
||||
}
|
||||
|
||||
// ?results=1 inlines every result.
|
||||
full := batchDecode(t, batchReq(t, s, http.MethodGet, "/v1/batch/"+id+"?results=1", keys["test"], nil))
|
||||
for i, x := range full["targets"].([]any) {
|
||||
if _, has := x.(map[string]any)["result"]; !has {
|
||||
t.Fatalf("results=1: target %d has no result", i)
|
||||
}
|
||||
}
|
||||
|
||||
if got := batchReq(t, s, http.MethodGet, "/v1/batch/"+id+"/2", keys["test"], nil).Code; got != 404 {
|
||||
t.Fatalf("GET idx out of range: %d, want 404", got)
|
||||
}
|
||||
if got := batchReq(t, s, http.MethodGet, "/v1/batch/"+id+"/x", keys["test"], nil).Code; got != 404 {
|
||||
t.Fatalf("GET idx not a number: %d, want 404", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchIdempotentResubmit(t *testing.T) {
|
||||
ts := pageServer(t, nil)
|
||||
s, keys := newBatchServer(t, "test")
|
||||
tg := map[string]any{"url": ts.URL, "mode": "clean"}
|
||||
a := submit(t, s, keys["test"], tg)
|
||||
b := submit(t, s, keys["test"], tg)
|
||||
if a["job_id"] != b["job_id"] || b["reused"] != true {
|
||||
t.Fatalf("resubmit: %v then %v", a, b)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchRefusedUpFront(t *testing.T) {
|
||||
s, keys := newBatchServer(t, "test")
|
||||
cases := []struct {
|
||||
name string
|
||||
body any
|
||||
code int
|
||||
want string
|
||||
}{
|
||||
{"guard pre-check (port)", map[string]any{"targets": []any{map[string]any{"url": "http://127.0.0.1/"}, map[string]any{"url": "http://example.com:8080/"}}}, 422, "blocked"},
|
||||
{"guard pre-check (private literal)", map[string]any{"targets": []any{map[string]any{"url": "http://10.0.0.5/"}}}, 422, "blocked"},
|
||||
{"no targets", map[string]any{"targets": []any{}}, 400, ""},
|
||||
{"not json", "nope", 400, ""},
|
||||
{"empty url", map[string]any{"targets": []any{map[string]any{"url": ""}}}, 400, ""},
|
||||
{"bad mode", map[string]any{"targets": []any{map[string]any{"url": "http://127.0.0.1/", "mode": "shout"}}}, 400, ""},
|
||||
}
|
||||
for _, c := range cases {
|
||||
rec := batchReq(t, s, http.MethodPost, "/v1/batch", keys["test"], c.body)
|
||||
if rec.Code != c.code || !strings.Contains(rec.Body.String(), c.want) {
|
||||
t.Errorf("%s: %d %s, want %d containing %q", c.name, rec.Code, rec.Body.String(), c.code, c.want)
|
||||
}
|
||||
}
|
||||
// The guard refusal names the offending index.
|
||||
rec := batchReq(t, s, http.MethodPost, "/v1/batch", keys["test"], cases[0].body)
|
||||
if !strings.Contains(rec.Body.String(), "1") {
|
||||
t.Errorf("guard refusal does not name target 1: %s", rec.Body.String())
|
||||
}
|
||||
|
||||
var many []any
|
||||
for i := 0; i < 33; i++ {
|
||||
many = append(many, map[string]any{"url": fmt.Sprintf("http://127.0.0.1/%d", i)})
|
||||
}
|
||||
if rec := batchReq(t, s, http.MethodPost, "/v1/batch", keys["test"], map[string]any{"targets": many}); rec.Code != 422 {
|
||||
t.Errorf("33 targets: %d, want 422", rec.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// 32 targets, each with a 500-rune question, fit the batch body cap even though they are well
|
||||
// over the 64 KiB cap of the other endpoints.
|
||||
func TestBatchBodyCap(t *testing.T) {
|
||||
hold := make(chan struct{})
|
||||
defer close(hold)
|
||||
ts := pageServer(t, hold)
|
||||
s, keys := newBatchServer(t, "test")
|
||||
q := strings.Repeat("q", 500)
|
||||
var many []any
|
||||
for i := 0; i < 32; i++ {
|
||||
many = append(many, map[string]any{"url": fmt.Sprintf("%s/%d?pad=%s", ts.URL, i, strings.Repeat("p", 1500)), "question": q})
|
||||
}
|
||||
if rec := batchReq(t, s, http.MethodPost, "/v1/batch", keys["test"], map[string]any{"targets": many}); rec.Code != 202 {
|
||||
t.Fatalf("32 targets with questions (~%d KiB): %d %s", 32*2100/1024, rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchPartialFailure(t *testing.T) {
|
||||
ts := pageServer(t, nil)
|
||||
s, keys := newBatchServer(t, "test")
|
||||
// Target 1 is a name that resolves to a private address. Names are not resolved at submit,
|
||||
// so it passes the pre-check; the fetch-time guard blocks it: a per-target failure.
|
||||
sub := submit(t, s, keys["test"], map[string]any{"url": ts.URL}, map[string]any{"url": "http://internal.test/"})
|
||||
env := waitTerminal(t, s, keys["test"], sub["job_id"].(string))
|
||||
if env["status"] != "partial" {
|
||||
t.Fatalf("status %v", env["status"])
|
||||
}
|
||||
second := env["targets"].([]any)[1].(map[string]any)
|
||||
if second["status"] != "failed" {
|
||||
t.Fatalf("target 1: %v", second)
|
||||
}
|
||||
if e, _ := second["error"].(string); !strings.Contains(e, "blocked") {
|
||||
t.Fatalf("target 1 error %q", e)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchOwnership404(t *testing.T) {
|
||||
hold := make(chan struct{})
|
||||
defer close(hold)
|
||||
ts := pageServer(t, hold)
|
||||
s, keys := newBatchServer(t, "test", "other")
|
||||
id := submit(t, s, keys["test"], map[string]any{"url": ts.URL})["job_id"].(string)
|
||||
for _, p := range []string{"/v1/batch/" + id, "/v1/batch/" + id + "/0"} {
|
||||
if got := batchReq(t, s, http.MethodGet, p, keys["other"], nil).Code; got != 404 {
|
||||
t.Fatalf("other caller GET %s: %d (must not reveal the job)", p, got)
|
||||
}
|
||||
}
|
||||
if got := batchReq(t, s, http.MethodDelete, "/v1/batch/"+id, keys["other"], nil).Code; got != 404 {
|
||||
t.Fatalf("other caller DELETE: %d", got)
|
||||
}
|
||||
if got := batchReq(t, s, http.MethodGet, "/v1/batch/"+id, keys["test"], nil).Code; got != 200 {
|
||||
t.Fatalf("owner GET: %d", got)
|
||||
}
|
||||
if got := batchReq(t, s, http.MethodGet, "/v1/batch/no-such-job", keys["test"], nil).Code; got != 404 {
|
||||
t.Fatalf("unknown job: %d", got)
|
||||
}
|
||||
if got := batchReq(t, s, http.MethodGet, "/v1/batch/"+id, "not-a-key", nil).Code; got != 401 {
|
||||
t.Fatalf("bad key: %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchCancel(t *testing.T) {
|
||||
hold := make(chan struct{})
|
||||
ts := pageServer(t, hold)
|
||||
s, keys := newBatchServer(t, "test")
|
||||
id := submit(t, s, keys["test"], map[string]any{"url": ts.URL + "/a"}, map[string]any{"url": ts.URL + "/b"})["job_id"].(string)
|
||||
rec := batchReq(t, s, http.MethodDelete, "/v1/batch/"+id, keys["test"], nil)
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("cancel: %d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if b := batchDecode(t, rec); b["success"] != true || b["status"] != "cancelled" {
|
||||
t.Fatalf("cancel body %v", b)
|
||||
}
|
||||
close(hold)
|
||||
env := waitTerminal(t, s, keys["test"], id)
|
||||
if env["status"] != "cancelled" {
|
||||
t.Fatalf("status %v", env["status"])
|
||||
}
|
||||
if tg := env["targets"].([]any)[1].(map[string]any); tg["status"] != "pending" {
|
||||
t.Fatalf("target 1 after cancel: %v", tg)
|
||||
}
|
||||
if got := batchReq(t, s, http.MethodDelete, "/v1/batch/"+id, keys["test"], nil).Code; got != 409 {
|
||||
t.Fatalf("cancelling a finished job: %d, want 409", got)
|
||||
}
|
||||
}
|
||||
|
||||
// The request metrics never get a job id as a label value.
|
||||
func TestBatchMetricsPathLabel(t *testing.T) {
|
||||
ts := pageServer(t, nil)
|
||||
s, keys := newBatchServer(t, "test")
|
||||
id := submit(t, s, keys["test"], map[string]any{"url": ts.URL})["job_id"].(string)
|
||||
waitTerminal(t, s, keys["test"], id)
|
||||
batchReq(t, s, http.MethodGet, "/v1/batch/"+id+"/0", keys["test"], nil)
|
||||
if strings.Contains(promText(t, s), id) {
|
||||
t.Fatal("a job id appears in /metrics")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,202 @@
|
||||
package api
|
||||
|
||||
// P6 task 06: the runner wired to fetchTarget and the metrics, with no HTTP yet. The helpers
|
||||
// here (newBatchServer, pageServer, fakeDNS) are shared with task 07's batch_test.go.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/netip"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/auth"
|
||||
"git.wntrmute.dev/kyle/sift/internal/batch"
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
"git.wntrmute.dev/kyle/sift/internal/fetch"
|
||||
"git.wntrmute.dev/kyle/sift/internal/guard"
|
||||
)
|
||||
|
||||
var loopback = []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}
|
||||
|
||||
// fakeDNS resolves internal.test to a private address; every other name fails to resolve.
|
||||
func fakeDNS(_ context.Context, host string) ([]netip.Addr, error) {
|
||||
if host == "internal.test" {
|
||||
return []netip.Addr{netip.MustParseAddr("10.0.0.5")}, nil
|
||||
}
|
||||
return nil, fmt.Errorf("no such host %s", host)
|
||||
}
|
||||
|
||||
// newBatchServer builds a Server with the batch runner wired over an in-memory dcache through
|
||||
// s.NewBatchRunner, exactly as serve does. callers are the names to provision (each gets a key).
|
||||
func newBatchServer(t *testing.T, callers ...string) (*Server, map[string]string) {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
keysPath := filepath.Join(dir, "keys.toml")
|
||||
keys := map[string]string{}
|
||||
for _, name := range callers {
|
||||
out := filepath.Join(dir, name)
|
||||
if err := auth.Add(keysPath, name, "hermes-fleet", out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
b, err := os.ReadFile(out)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
keys[name] = strings.TrimSpace(string(b))
|
||||
}
|
||||
store, err := auth.Load(keysPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
g := &guard.Guard{AllowPrefixes: loopback, LookupIP: fakeDNS}
|
||||
f := &fetch.Fetcher{Guard: g, UA: "sift-test"}
|
||||
log := slog.New(slog.NewTextHandler(io.Discard, nil))
|
||||
s := New("test", g, f, store, log)
|
||||
db, err := dcache.Open("", 1<<20, log)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
s.Batch = s.NewBatchRunner(db, batch.Config{})
|
||||
t.Cleanup(s.Batch.Close) // before db.Close
|
||||
return s, keys
|
||||
}
|
||||
|
||||
// pageServer serves a small HTML page at every path; hold, when non-nil, blocks each request
|
||||
// until it is closed.
|
||||
func pageServer(t *testing.T, hold chan struct{}) *httptest.Server {
|
||||
t.Helper()
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if hold != nil {
|
||||
select {
|
||||
case <-hold:
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
w.Header().Set("Content-Type", "text/html; charset=utf-8")
|
||||
fmt.Fprintf(w, `<html><head><title>T</title></head><body><article><p>hello batch %s</p></article></body></html>`, r.URL.Path)
|
||||
}))
|
||||
t.Cleanup(ts.Close)
|
||||
return ts
|
||||
}
|
||||
|
||||
func waitJob(t *testing.T, r *batch.Runner, id string) batch.Job {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(10 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
j, _, ok := r.Status(id)
|
||||
if ok && (j.Status == "done" || j.Status == "partial" || j.Status == "failed" || j.Status == "cancelled") {
|
||||
return j
|
||||
}
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
t.Fatal("job never finished")
|
||||
return batch.Job{}
|
||||
}
|
||||
|
||||
func promText(t *testing.T, s *Server) string {
|
||||
t.Helper()
|
||||
req := httptest.NewRequest(http.MethodGet, "/metrics", nil)
|
||||
req.RemoteAddr, req.Host = "127.0.0.1:9", "127.0.0.1:8890"
|
||||
rec := httptest.NewRecorder()
|
||||
s.ServeHTTP(rec, req)
|
||||
return rec.Body.String()
|
||||
}
|
||||
|
||||
// The runner fetches through fetchTarget as the submitting caller; a name that resolves to a
|
||||
// private address is a per-target "blocked" failure; the hooks reach the metrics.
|
||||
func TestBatchRunnerWired(t *testing.T) {
|
||||
ts := pageServer(t, nil)
|
||||
s, _ := newBatchServer(t, "test")
|
||||
c := &auth.Caller{Name: "test", Workspace: "hermes-fleet"}
|
||||
id, reused, err := s.Batch.Submit(context.Background(), c, dcache.ClassFleet, []batch.Target{
|
||||
{URL: ts.URL + "/one"}, {URL: "http://internal.test/"},
|
||||
})
|
||||
if err != nil || reused {
|
||||
t.Fatalf("submit: %v %v", err, reused)
|
||||
}
|
||||
if j := waitJob(t, s.Batch, id); j.Status != "partial" {
|
||||
t.Fatalf("job %+v", j)
|
||||
}
|
||||
_, ts2, _ := s.Batch.Status(id)
|
||||
if ts2[0].Status != "done" || ts2[0].Result == nil || !strings.Contains(ts2[0].Result.Content, "hello batch /one") {
|
||||
t.Fatalf("target 0: %+v", ts2[0])
|
||||
}
|
||||
if ts2[0].Result.Meta["mode"] != "clean" || ts2[0].Result.Meta["untrusted"] != true {
|
||||
t.Fatalf("target 0 meta %v: not the fetchTarget envelope", ts2[0].Result.Meta)
|
||||
}
|
||||
if ts2[1].Status != "failed" || !strings.Contains(ts2[1].Error, "blocked") {
|
||||
t.Fatalf("target 1: %+v", ts2[1])
|
||||
}
|
||||
out := promText(t, s)
|
||||
for _, want := range []string{
|
||||
`sift_batch_jobs_total{outcome="created"} 1`,
|
||||
`sift_batch_jobs_total{outcome="partial"} 1`,
|
||||
`sift_batch_targets_total{outcome="ok"} 1`,
|
||||
`sift_batch_targets_total{outcome="blocked"} 1`,
|
||||
`sift_batch_result_bytes_count 1`,
|
||||
`sift_batch_queue_depth 0`,
|
||||
} {
|
||||
if !strings.Contains(out, want) {
|
||||
t.Errorf("metrics lack %q", want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The adapter turns fetch errors into what the runner and the metrics need.
|
||||
func TestBatchOutcomeClasses(t *testing.T) {
|
||||
cases := []struct {
|
||||
err error
|
||||
want string
|
||||
}{
|
||||
{nil, "ok"},
|
||||
{fmt.Errorf("busy: %w", batch.ErrRetryLater), "busy"},
|
||||
{&batch.TargetError{Status: 403, Msg: "blocked: private address"}, "blocked"},
|
||||
{&batch.TargetError{Status: 429, Msg: "busy"}, "busy"},
|
||||
{&batch.TargetError{Status: 504, Msg: "timeout"}, "timeout"},
|
||||
{&batch.TargetError{Status: 415, Msg: "unsupported content type"}, "unsupported"},
|
||||
{&batch.TargetError{Status: 422, Msg: "document too long for question (6001 chunks); use grep or pages"}, "too_long"},
|
||||
{&batch.TargetError{Status: 422, Msg: "unreadable PDF"}, "upstream_4xx"},
|
||||
{&batch.TargetError{Status: 404, Msg: "not found"}, "upstream_4xx"},
|
||||
{&batch.TargetError{Status: 500, Msg: "internal error"}, "upstream_5xx"},
|
||||
{errors.New("something else"), "upstream_5xx"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := batchOutcome(c.err); got != c.want {
|
||||
t.Errorf("batchOutcome(%v) = %q, want %q", c.err, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// batchFetch maps a busy fetch to ErrRetryLater (the runner retries it) and every other
|
||||
// failure to a TargetError with the /v1/fetch status.
|
||||
func TestBatchFetchErrors(t *testing.T) {
|
||||
s, _ := newBatchServer(t, "test")
|
||||
c := &auth.Caller{Name: "test", Workspace: "hermes-fleet"}
|
||||
_, err := s.batchFetch(context.Background(), c, batch.Target{URL: "http://internal.test/"})
|
||||
var te *batch.TargetError
|
||||
if !errors.As(err, &te) || te.Status != 403 || !strings.Contains(te.Msg, "blocked") {
|
||||
t.Fatalf("private name: %v, want a TargetError 403 blocked", err)
|
||||
}
|
||||
// Fill the global fetch gate so the next fetch is refused busy.
|
||||
for i := 0; i < 8; i++ {
|
||||
if err := s.acquire(fmt.Sprintf("filler%d", i), "http://filler.test/"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
ts := pageServer(t, nil)
|
||||
_, err = s.batchFetch(context.Background(), c, batch.Target{URL: ts.URL})
|
||||
if !errors.Is(err, batch.ErrRetryLater) {
|
||||
t.Fatalf("busy gate: %v, want an error wrapping batch.ErrRetryLater", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,220 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/clean"
|
||||
"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/pdfdoc"
|
||||
"git.wntrmute.dev/kyle/sift/internal/pdftext"
|
||||
)
|
||||
|
||||
// v1extract answers a PDF the caller POSTs in the request body, exactly as /v1/fetch
|
||||
// answers the same PDF by URL. Nothing is fetched: the bytes are hashed, the hash is the
|
||||
// cache key, and identical content is extracted once (content-addressed, so no TTL and no
|
||||
// revalidation).
|
||||
func (s *Server) v1extract(w http.ResponseWriter, r *http.Request) {
|
||||
s.reqN.Add(1)
|
||||
c := s.caller(w, r)
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case s.uploads <- struct{}{}:
|
||||
defer func() { <-s.uploads }()
|
||||
case <-time.After(uploadWait):
|
||||
w.Header().Set("Retry-After", "10")
|
||||
writeErr(w, 429, "busy", "")
|
||||
return
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
}
|
||||
|
||||
body, rerr := io.ReadAll(r.Body)
|
||||
if rerr != nil {
|
||||
var me *http.MaxBytesError
|
||||
if errors.As(rerr, &me) {
|
||||
writeErr(w, 413, fmt.Sprintf("too large (>%d MB)", s.pdfMaxUpload()>>20), "")
|
||||
return
|
||||
}
|
||||
writeErr(w, 400, "empty body", "")
|
||||
return
|
||||
}
|
||||
if len(body) == 0 {
|
||||
writeErr(w, 400, "empty body", "")
|
||||
return
|
||||
}
|
||||
if clean.Kind(r.Header.Get("Content-Type"), body) != "pdf" {
|
||||
writeErr(w, 415, "extract accepts PDFs only", "")
|
||||
return
|
||||
}
|
||||
|
||||
q := r.URL.Query()
|
||||
pages := q.Get("pages")
|
||||
grep := q.Get("grep")
|
||||
question := q.Get("question")
|
||||
mode := q.Get("mode")
|
||||
if mode == "" {
|
||||
if question != "" {
|
||||
mode = "digest"
|
||||
} else {
|
||||
mode = "clean"
|
||||
}
|
||||
}
|
||||
if mode != "clean" && mode != "digest" {
|
||||
writeErr(w, 400, "invalid mode", "")
|
||||
return
|
||||
}
|
||||
|
||||
hash := "sha256:" + HashBody(body)
|
||||
key := dcache.PDFKey("upload:"+hash, s.callerClass(c))
|
||||
|
||||
t0 := time.Now()
|
||||
var res *fetch.Resource
|
||||
var doc *pdftext.Doc
|
||||
cached := false
|
||||
if cache := s.cache(); cache != nil {
|
||||
if e, ok := cache.GetPDF(key); ok {
|
||||
res, doc = s.pdfFromEntry("upload", e)
|
||||
cached = true
|
||||
s.Prom.PDFCacheHit()
|
||||
}
|
||||
}
|
||||
if res == nil {
|
||||
res = &fetch.Resource{URL: "upload", FinalURL: "upload", Status: 200, ContentType: "application/pdf", Body: body, Hash: hash, FetchedAt: t0}
|
||||
var eerr error
|
||||
doc, eerr = s.extractPDF(r.Context(), c, res)
|
||||
if eerr != nil {
|
||||
s.writeFetchErr(w, c, "", eerr)
|
||||
return
|
||||
}
|
||||
s.storePDF(key, res, doc)
|
||||
}
|
||||
fetchMS := time.Since(t0).Milliseconds()
|
||||
|
||||
rp, err := s.renderPDF(doc, pdfdoc.Options{Pages: pages, Grep: grep, OutCap: s.OutCap})
|
||||
if err != nil {
|
||||
if errors.Is(err, pdfdoc.ErrBadPages) || errors.Is(err, pdfdoc.ErrBadGrep) {
|
||||
s.writeFetchErr(w, c, "", &fetch.Error{Status: 400, Msg: err.Error()})
|
||||
return
|
||||
}
|
||||
s.writeFetchErr(w, c, "", err)
|
||||
return
|
||||
}
|
||||
full := pdfdoc.Full(doc)
|
||||
docLen := utf8.RuneCountInString(full)
|
||||
content := rp.Content
|
||||
if mode == "digest" {
|
||||
content = full
|
||||
}
|
||||
|
||||
warn := []string{}
|
||||
var digMeta map[string]any
|
||||
var digestMS int64
|
||||
if mode == "digest" {
|
||||
n := len(digest.Split(full, digest.ChunkTokens))
|
||||
limit := s.PDFMaxQuestionChunks
|
||||
if limit == 0 {
|
||||
limit = DefaultPDFMaxQuestionChunks
|
||||
}
|
||||
if n > limit {
|
||||
s.writeFetchErr(w, c, "", &fetch.Error{Status: 422, Msg: fmt.Sprintf("document too long for question (%d chunks); use grep or pages", n)})
|
||||
return
|
||||
}
|
||||
if s.Digest == nil {
|
||||
warn = append(warn, "digest fell back to clean: unset")
|
||||
mode = "clean"
|
||||
content, warn = capDigestFallback(content, warn)
|
||||
digMeta = map[string]any{"fallback": "unset"}
|
||||
} else {
|
||||
dr := s.Digest.Run(r.Context(), content, question, res.FinalURL, s.callerClass(c))
|
||||
digestMS = dr.DigestMS
|
||||
if !dr.Cached {
|
||||
s.Prom.ObserveDigest(time.Duration(dr.DigestMS) * time.Millisecond)
|
||||
}
|
||||
if dr.Fallback != "" {
|
||||
warn = append(warn, "digest fell back to clean: "+dr.Fallback)
|
||||
mode = "clean"
|
||||
content, warn = capDigestFallback(content, warn)
|
||||
digMeta = map[string]any{"fallback": dr.Fallback, "model": dr.Model}
|
||||
} else {
|
||||
content = dr.Text
|
||||
digMeta = map[string]any{
|
||||
"model": dr.Model,
|
||||
"chunks_used": dr.ChunksUsed,
|
||||
"chunks_total": dr.ChunksTotal,
|
||||
"retrieval": dr.Retrieval,
|
||||
"embedded": dr.Embedded,
|
||||
"fallback": false,
|
||||
"cached": dr.Cached,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pdfMap := map[string]any{
|
||||
"pages": doc.Pages,
|
||||
"title": doc.Title,
|
||||
"outline": doc.Outline,
|
||||
"extracted": len(doc.Text),
|
||||
"first": rp.First,
|
||||
"last": rp.Last,
|
||||
"map": rp.Map,
|
||||
"grep_hits": rp.GrepHits,
|
||||
"note": rp.Note,
|
||||
}
|
||||
|
||||
s.Log.Info("extract", "caller", c.Name, "bytes", len(body), "pdf_pages", doc.Pages, "mode", mode, "cached", cached)
|
||||
|
||||
src := map[string]any{
|
||||
"kind": "upload",
|
||||
"hash": hash,
|
||||
"bytes": len(body),
|
||||
"cached": cached,
|
||||
"fetched_at": res.FetchedAt.Format(time.RFC3339),
|
||||
}
|
||||
out := map[string]any{
|
||||
"content": content,
|
||||
"mode": mode,
|
||||
"untrusted": true,
|
||||
"source": src,
|
||||
"warnings": warn,
|
||||
"sizes": map[string]int{
|
||||
"raw_tokens": docLen / 4,
|
||||
"clean_tokens": digest.EstTokens(full),
|
||||
"returned_tokens": digest.EstTokens(content),
|
||||
},
|
||||
"timings": map[string]int64{"fetch_ms": fetchMS, "digest_ms": digestMS},
|
||||
}
|
||||
if pdfMap != nil {
|
||||
out["pdf"] = pdfMap
|
||||
}
|
||||
if digMeta != nil {
|
||||
out["digest"] = digMeta
|
||||
}
|
||||
writeJSON(w, 200, out)
|
||||
}
|
||||
|
||||
// maxUploads is how many uploads may be in memory at once, and uploadWait how long a further
|
||||
// one waits for a slot before a 429: at 100 MiB each, a few are all a 1 GiB service can hold
|
||||
// beside the extraction children.
|
||||
const (
|
||||
maxUploads = 2
|
||||
uploadWait = 5 * time.Second
|
||||
)
|
||||
|
||||
// pdfMaxUpload is the /v1/extract upload cap: the configured value, or the default when unset.
|
||||
func (s *Server) pdfMaxUpload() int64 {
|
||||
if s.PDFMaxUpload > 0 {
|
||||
return s.PDFMaxUpload
|
||||
}
|
||||
return DefaultPDFMaxUpload
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
package api
|
||||
|
||||
// P5c owner review: regression tests for /v1/extract.
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Uploads are buffered whole (up to 100 MiB) before the PDF pool sees them, so the number read
|
||||
// at once is bounded: a third concurrent upload waits, then is refused busy.
|
||||
func TestExtractUploadsBounded(t *testing.T) {
|
||||
r := newPDFRig(t)
|
||||
for i := 0; i < cap(r.s.uploads); i++ {
|
||||
r.s.uploads <- struct{}{}
|
||||
}
|
||||
start := time.Now()
|
||||
code, out := r.extract(t, r.fleet, "", threePDF(), "application/pdf")
|
||||
if code != http.StatusTooManyRequests {
|
||||
t.Fatalf("upload with every slot taken: %d %v, want 429", code, out)
|
||||
}
|
||||
if d := time.Since(start); d < uploadWait-100*time.Millisecond {
|
||||
t.Fatalf("refused after %v, want a wait of about %v first", d, uploadWait)
|
||||
}
|
||||
<-r.s.uploads
|
||||
if code, out := r.extract(t, r.fleet, "", threePDF(), "application/pdf"); code != 200 {
|
||||
t.Fatalf("upload with a free slot: %d %v", code, out)
|
||||
}
|
||||
}
|
||||
|
||||
// The envelope's warnings are a list, never null, as in /v1/fetch.
|
||||
func TestExtractWarningsIsAList(t *testing.T) {
|
||||
r := newPDFRig(t)
|
||||
code, out := r.extract(t, r.fleet, "", threePDF(), "application/pdf")
|
||||
if code != 200 {
|
||||
t.Fatalf("%d %v", code, out)
|
||||
}
|
||||
if _, ok := out["warnings"].([]any); !ok {
|
||||
t.Fatalf("warnings = %#v, want a list", out["warnings"])
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,122 @@
|
||||
package api
|
||||
|
||||
// P5c task 01: POST /v1/extract takes a PDF's bytes (a file already on the caller's host), and
|
||||
// answers exactly like /v1/fetch does for a PDF URL. The cache key is the content hash.
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func (r *pdfRig) extract(t *testing.T, key, query string, body []byte, ct string) (int, map[string]any) {
|
||||
t.Helper()
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/extract"+query, bytes.NewReader(body))
|
||||
if ct != "" {
|
||||
req.Header.Set("Content-Type", ct)
|
||||
}
|
||||
if key != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
}
|
||||
rec := httptest.NewRecorder()
|
||||
r.s.ServeHTTP(rec, req)
|
||||
var out map[string]any
|
||||
_ = json.Unmarshal(rec.Body.Bytes(), &out)
|
||||
return rec.Code, out
|
||||
}
|
||||
|
||||
func TestExtractUpload(t *testing.T) {
|
||||
r := newPDFRig(t)
|
||||
pdf := threePDF()
|
||||
code, out := r.extract(t, r.fleet, "", pdf, "application/pdf")
|
||||
if code != 200 {
|
||||
t.Fatalf("%d %v", code, out)
|
||||
}
|
||||
content, _ := out["content"].(string)
|
||||
if !strings.Contains(content, "--- page 2 ---") || !strings.Contains(content, "quick brown fox") {
|
||||
t.Errorf("content %q", content)
|
||||
}
|
||||
if p := pdfInfo(t, out); p["pages"] != 3.0 || p["title"] != "Three Pages" {
|
||||
t.Errorf("pdf %v", p)
|
||||
}
|
||||
src, _ := out["source"].(map[string]any)
|
||||
if src["kind"] != "upload" || src["bytes"] != float64(len(pdf)) || !strings.HasPrefix(fmt.Sprint(src["hash"]), "sha256:") || src["cached"] != false {
|
||||
t.Errorf("source %v", src)
|
||||
}
|
||||
if _, hasURL := src["url"]; hasURL {
|
||||
t.Error("an upload has no URL; source must not invent one")
|
||||
}
|
||||
// The same bytes again: from the cache, no child.
|
||||
code, out = r.extract(t, r.fleet, "?pages=2", pdf, "application/octet-stream")
|
||||
if code != 200 || r.runs() != 1 {
|
||||
t.Fatalf("second upload: %d, %d extractions; want 200 and 1", code, r.runs())
|
||||
}
|
||||
if src, _ := out["source"].(map[string]any); src["cached"] != true {
|
||||
t.Errorf("second upload source %v, want cached", src)
|
||||
}
|
||||
if c, _ := out["content"].(string); strings.Contains(c, "--- page 1 ---") || !strings.Contains(c, "--- page 2 ---") {
|
||||
t.Errorf("pages=2: %q", c)
|
||||
}
|
||||
code, out = r.extract(t, r.fleet, "?grep=quick+brown", pdf, "application/pdf")
|
||||
if c, _ := out["content"].(string); code != 200 || !strings.Contains(c, "p. 2: The quick brown fox") {
|
||||
t.Errorf("grep: %d %q", code, c)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractRefusals(t *testing.T) {
|
||||
r := newPDFRig(t)
|
||||
if code, _ := r.extract(t, "", "", threePDF(), "application/pdf"); code != 401 {
|
||||
t.Errorf("no key: %d, want 401", code)
|
||||
}
|
||||
if code, out := r.extract(t, r.fleet, "", []byte("<html>not a pdf</html>"), "text/html"); code != 415 {
|
||||
t.Errorf("html body: %d %v, want 415", code, out)
|
||||
}
|
||||
if code, _ := r.extract(t, r.fleet, "", nil, "application/pdf"); code != 400 {
|
||||
t.Errorf("empty body: %d, want 400", code)
|
||||
}
|
||||
if code, _ := r.extract(t, r.fleet, "?pages=0-1", threePDF(), "application/pdf"); code != 400 {
|
||||
t.Errorf("bad pages: %d, want 400", code)
|
||||
}
|
||||
r.s.PDFMaxUpload = 1000
|
||||
if code, out := r.extract(t, r.fleet, "", threePDF(), "application/pdf"); code != 413 {
|
||||
t.Errorf("over PDFMaxUpload: %d %v, want 413", code, out)
|
||||
}
|
||||
if DefaultPDFMaxUpload != 100<<20 {
|
||||
t.Errorf("DefaultPDFMaxUpload = %d, want 100 MiB", DefaultPDFMaxUpload)
|
||||
}
|
||||
// The ordinary JSON endpoints keep their small request limit.
|
||||
big := bytes.Repeat([]byte("x"), 128<<10)
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/fetch", bytes.NewReader(append([]byte(`{"url":"`), append(big, []byte(`"}`)...)...)))
|
||||
req.Header.Set("Authorization", "Bearer "+r.fleet)
|
||||
r.s.ServeHTTP(rec, req)
|
||||
if rec.Code == 200 {
|
||||
t.Error("a 128 KiB /v1/fetch body was accepted; the JSON limit must still apply there")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractRestrictedPartition(t *testing.T) {
|
||||
r := newPDFRig(t)
|
||||
pdf := threePDF()
|
||||
if code, _ := r.extract(t, r.gs, "", pdf, "application/pdf"); code != 200 {
|
||||
t.Fatalf("restricted upload: %d", code)
|
||||
}
|
||||
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 })
|
||||
if strings.Contains(b.String(), "Three Pages") || strings.Contains(b.String(), "quick brown") {
|
||||
t.Errorf("an upload's log line carries document text: %q", b.String())
|
||||
}
|
||||
}
|
||||
r.logs.mu.Unlock()
|
||||
if code, _ := r.extract(t, r.fleet, "", pdf, "application/pdf"); code != 200 || r.runs() != 2 {
|
||||
t.Errorf("fleet after restricted: %d, %d extractions; want 200 and 2 (separate partitions)", code, r.runs())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
package api
|
||||
|
||||
// P6 task 05: v1fetch's per-target pipeline is extracted into s.fetchTarget, which the batch
|
||||
// runner also calls. These tests pin that the two agree; every existing v1fetch test must also
|
||||
// pass unchanged.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/netip"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/auth"
|
||||
"git.wntrmute.dev/kyle/sift/internal/batch"
|
||||
"git.wntrmute.dev/kyle/sift/internal/fetch"
|
||||
"git.wntrmute.dev/kyle/sift/internal/guard"
|
||||
)
|
||||
|
||||
// newTargetServer is setupAPI with loopback allowed; its key belongs to caller "test".
|
||||
func newTargetServer(t *testing.T) (*Server, string) {
|
||||
t.Helper()
|
||||
return setupAPI(t, &guard.Guard{AllowPrefixes: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")}})
|
||||
}
|
||||
|
||||
func TestFetchTargetMatchesV1Fetch(t *testing.T) {
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/html; charset=utf-8")
|
||||
_, _ = w.Write([]byte(`<html><head><title>Same</title></head><body><article><h1>Same</h1><p>one pipeline, two callers</p></article></body></html>`))
|
||||
}))
|
||||
defer ts.Close()
|
||||
s, key := newTargetServer(t)
|
||||
c := &auth.Caller{Name: "test", Workspace: "hermes-fleet"}
|
||||
|
||||
res, err := s.fetchTarget(context.Background(), c, batch.Target{URL: ts.URL})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(res.Content, "one pipeline, two callers") {
|
||||
t.Fatalf("content %q", res.Content)
|
||||
}
|
||||
if _, has := res.Meta["content"]; has {
|
||||
t.Fatal("Meta must not repeat the content")
|
||||
}
|
||||
for _, k := range []string{"mode", "untrusted", "source", "warnings", "sizes", "timings"} {
|
||||
if _, has := res.Meta[k]; !has {
|
||||
t.Errorf("Meta lacks %q (it carries every envelope field except content)", k)
|
||||
}
|
||||
}
|
||||
if res.Meta["mode"] != "clean" || res.Meta["untrusted"] != true {
|
||||
t.Fatalf("meta %v", res.Meta)
|
||||
}
|
||||
|
||||
// The same URL through /v1/fetch: identical content and mode, and the envelope has exactly
|
||||
// the fields of Meta plus content.
|
||||
rec := postJSON(t, s, "/v1/fetch", key, map[string]any{"url": ts.URL})
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("v1fetch %d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var env map[string]any
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &env); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if env["content"] != res.Content || env["mode"] != res.Meta["mode"] {
|
||||
t.Fatalf("v1fetch and fetchTarget disagree:\n%v\n%+v", env, res)
|
||||
}
|
||||
if len(env) != len(res.Meta)+1 {
|
||||
t.Fatalf("envelope keys %d, Meta keys %d: want Meta + content", len(env), len(res.Meta))
|
||||
}
|
||||
}
|
||||
|
||||
func TestFetchTargetErrors(t *testing.T) {
|
||||
s, _ := newTargetServer(t)
|
||||
c := &auth.Caller{Name: "test", Workspace: "hermes-fleet"}
|
||||
cases := []struct {
|
||||
name string
|
||||
target batch.Target
|
||||
status int
|
||||
msg string
|
||||
}{
|
||||
{"blocked", batch.Target{URL: "http://10.0.0.5/"}, 403, "blocked"},
|
||||
{"bad mode", batch.Target{URL: "http://127.0.0.1/", Options: batch.Options{Mode: "shout"}}, 400, ""},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
_, err := s.fetchTarget(context.Background(), c, tc.target)
|
||||
var fe *fetch.Error
|
||||
if !errors.As(err, &fe) || fe.Status != tc.status || !strings.Contains(fe.Msg, tc.msg) {
|
||||
t.Errorf("%s: err %v, want a *fetch.Error %d containing %q", tc.name, err, tc.status, tc.msg)
|
||||
}
|
||||
}
|
||||
}
|
||||
+11
-4
@@ -18,10 +18,13 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
// DefaultPDFTTL bounds a cached PDF's life. Task 07; declared now so the test compiles.
|
||||
// DefaultPDFTTL is how long a cached PDF is served before it is revalidated.
|
||||
DefaultPDFTTL = 24 * time.Hour
|
||||
// DefaultPDFMaxQuestionChunks caps a digest question's length. Task 07; declared now.
|
||||
DefaultPDFMaxQuestionChunks = 200
|
||||
// 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
|
||||
@@ -55,7 +58,11 @@ func (s *Server) extractPDF(ctx context.Context, c *auth.Caller, res *fetch.Reso
|
||||
return v, err
|
||||
}
|
||||
start := time.Now()
|
||||
v, err := s.pool().Do(ctx, s.callerClass(c)+"|"+res.SHA256(), c.Name, fn)
|
||||
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.
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
package api
|
||||
|
||||
// P5b task 03: a question on a long PDF is answered (by BM25 first) instead of refused, and the
|
||||
// envelope says how the digest was built.
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/digest"
|
||||
"git.wntrmute.dev/kyle/sift/internal/pdftest"
|
||||
)
|
||||
|
||||
// fakeModel answers /embed with a fixed vector and /chat with a one-word SSE reply.
|
||||
func fakeModel(t *testing.T) *httptest.Server {
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/embed", func(w http.ResponseWriter, r *http.Request) {
|
||||
var body struct {
|
||||
Input []string `json:"input"`
|
||||
}
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
var data []map[string]any
|
||||
for i := range body.Input {
|
||||
data = append(data, map[string]any{"index": i, "embedding": []float64{1, float64(i % 7)}})
|
||||
}
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"data": data})
|
||||
})
|
||||
mux.HandleFunc("/chat", func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = w.Write([]byte(`data: {"choices":[{"delta":{"content":"answer (p. 3)"}}]}` + "\n\ndata: [DONE]\n\n"))
|
||||
})
|
||||
srv := httptest.NewServer(mux)
|
||||
t.Cleanup(srv.Close)
|
||||
return srv
|
||||
}
|
||||
|
||||
func TestPDFLongQuestionIsAnswered(t *testing.T) {
|
||||
r := newPDFRig(t)
|
||||
m := fakeModel(t)
|
||||
db := r.s.Digest.Cache
|
||||
r.s.Digest = &digest.Engine{ChatURL: m.URL + "/chat", EmbedURL: m.URL + "/embed", Cache: db, Gate: digest.NewGate(2, 8, 5*time.Second), Client: m.Client()}
|
||||
var pages []string
|
||||
for p := 1; p <= 260; p++ { // ~260 pages × 40 lines: well over EmbedSyncMax chunks
|
||||
pages = append(pages, pdftest.Lorem(p, 40))
|
||||
}
|
||||
o := newOrigin(t, "application/pdf", pdftest.Build(pdftest.Doc{Title: "Long", Pages: pages}))
|
||||
code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL, "question": "what is on page three?", "context": false})
|
||||
if code != 200 {
|
||||
t.Fatalf("long question: %d %v, want 200 (P5a refused this)", code, out)
|
||||
}
|
||||
dg, _ := out["digest"].(map[string]any)
|
||||
if dg == nil || dg["retrieval"] != "bm25" || dg["embedded"] != 0.0 {
|
||||
t.Errorf("digest meta = %v, want retrieval bm25 and embedded 0 on the first question", dg)
|
||||
}
|
||||
if c, _ := out["content"].(string); !strings.Contains(c, "answer") {
|
||||
t.Errorf("content %q", c)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPDFSmallQuestionReportsRetrieval(t *testing.T) {
|
||||
r := newPDFRig(t)
|
||||
m := fakeModel(t)
|
||||
db := r.s.Digest.Cache
|
||||
r.s.Digest = &digest.Engine{ChatURL: m.URL + "/chat", EmbedURL: m.URL + "/embed", Cache: db, Gate: digest.NewGate(2, 8, 5*time.Second), Client: m.Client()}
|
||||
o := newOrigin(t, "application/pdf", threePDF())
|
||||
code, out, _ := r.fetch(t, r.fleet, map[string]any{"url": o.srv.URL, "question": "what is on page two?", "context": false})
|
||||
dg, _ := out["digest"].(map[string]any)
|
||||
if code != 200 || dg == nil || dg["retrieval"] != "whole" {
|
||||
t.Errorf("small document: %d digest %v, want retrieval whole", code, dg)
|
||||
}
|
||||
}
|
||||
@@ -396,7 +396,8 @@ func TestPDFQuestionTooLong(t *testing.T) {
|
||||
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)
|
||||
}
|
||||
if DefaultPDFMaxQuestionChunks != 200 {
|
||||
t.Errorf("DefaultPDFMaxQuestionChunks = %d, want 200", DefaultPDFMaxQuestionChunks)
|
||||
// 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)
|
||||
}
|
||||
}
|
||||
|
||||
+50
-10
@@ -11,12 +11,22 @@ import (
|
||||
|
||||
type v1SearchReq struct {
|
||||
Query string `json:"query"`
|
||||
Queries []string `json:"queries"`
|
||||
Question string `json:"question"`
|
||||
K int `json:"k"`
|
||||
Rerank *bool `json:"rerank"`
|
||||
Digest bool `json:"digest"`
|
||||
URL string `json:"url"`
|
||||
Limit int `json:"limit"`
|
||||
IncludeDomains []string `json:"include_domains"`
|
||||
ExcludeDomains []string `json:"exclude_domains"`
|
||||
IncludeDomainsMode string `json:"include_domains_mode"`
|
||||
TimeRange string `json:"time_range"`
|
||||
StartDate string `json:"start_date"`
|
||||
EndDate string `json:"end_date"`
|
||||
Topic string `json:"topic"`
|
||||
ExactMatch bool `json:"exact_match"`
|
||||
MinScore *float64 `json:"min_score"`
|
||||
}
|
||||
|
||||
type v2SearchReq struct {
|
||||
@@ -42,9 +52,11 @@ func (s *Server) handleSearch(w http.ResponseWriter, r *http.Request, v2 bool) {
|
||||
}
|
||||
dec := json.NewDecoder(r.Body)
|
||||
query := ""
|
||||
var queries []string
|
||||
k := 0
|
||||
doRerank := true
|
||||
digest := false
|
||||
var f search.Filter
|
||||
if v2 {
|
||||
var req v2SearchReq
|
||||
if err := dec.Decode(&req); err != nil && err.Error() != "EOF" {
|
||||
@@ -55,13 +67,17 @@ func (s *Server) handleSearch(w http.ResponseWriter, r *http.Request, v2 bool) {
|
||||
k = req.Limit
|
||||
digest = req.Digest
|
||||
_ = req.URL
|
||||
if msg := search.ValidateQuery(query); msg != "" {
|
||||
writeErr(w, 400, msg, "")
|
||||
return
|
||||
}
|
||||
queries = []string{strings.TrimSpace(query)}
|
||||
} else {
|
||||
var req v1SearchReq
|
||||
if err := dec.Decode(&req); err != nil && err.Error() != "EOF" {
|
||||
writeErr(w, 400, "invalid query", "")
|
||||
return
|
||||
}
|
||||
query = req.Query
|
||||
k = req.K
|
||||
if k == 0 {
|
||||
k = req.Limit
|
||||
@@ -72,11 +88,26 @@ func (s *Server) handleSearch(w http.ResponseWriter, r *http.Request, v2 bool) {
|
||||
digest = req.Digest
|
||||
_ = req.URL
|
||||
_ = req.Question
|
||||
raw := search.Filter{
|
||||
IncludeDomains: req.IncludeDomains,
|
||||
ExcludeDomains: req.ExcludeDomains,
|
||||
IncludeDomainsMode: req.IncludeDomainsMode,
|
||||
TimeRange: req.TimeRange,
|
||||
StartDate: req.StartDate,
|
||||
EndDate: req.EndDate,
|
||||
Topic: req.Topic,
|
||||
ExactMatch: req.ExactMatch,
|
||||
}
|
||||
if msg := search.ValidateQuery(query); msg != "" {
|
||||
writeErr(w, 400, msg, "")
|
||||
if req.MinScore != nil {
|
||||
raw.MinScore = *req.MinScore
|
||||
}
|
||||
var errMsg, hint string
|
||||
queries, f, errMsg, hint = search.NormalizeSearch(req.Query, req.Queries, raw)
|
||||
if errMsg != "" {
|
||||
writeErr(w, 400, errMsg, hint)
|
||||
return
|
||||
}
|
||||
}
|
||||
defK := search.DefaultK
|
||||
if v2 {
|
||||
defK = search.DefaultV2K
|
||||
@@ -86,7 +117,7 @@ func (s *Server) handleSearch(w http.ResponseWriter, r *http.Request, v2 bool) {
|
||||
writeErr(w, 422, "no search results", "")
|
||||
return
|
||||
}
|
||||
res := s.Search.Run(r.Context(), strings.TrimSpace(query), k, doRerank, c.Name, s.callerClass(c))
|
||||
res := s.Search.RunFilter(r.Context(), queries, k, doRerank, c.Name, s.callerClass(c), f)
|
||||
warns := append([]string(nil), res.Warnings...)
|
||||
if digest {
|
||||
warns = append(warns, "digest ignored")
|
||||
@@ -95,7 +126,7 @@ func (s *Server) handleSearch(w http.ResponseWriter, r *http.Request, v2 bool) {
|
||||
writeErr(w, 422, "no search results", "")
|
||||
return
|
||||
}
|
||||
s.logSearch(c, query, res)
|
||||
s.logSearch(c, queries, f, res)
|
||||
if v2 {
|
||||
web := make([]map[string]any, 0, len(res.Hits))
|
||||
for i, h := range res.Hits {
|
||||
@@ -111,13 +142,17 @@ func (s *Server) handleSearch(w http.ResponseWriter, r *http.Request, v2 bool) {
|
||||
}
|
||||
results := make([]map[string]any, 0, len(res.Hits))
|
||||
for _, h := range res.Hits {
|
||||
results = append(results, map[string]any{
|
||||
row := map[string]any{
|
||||
"url": h.URL,
|
||||
"title": h.Title,
|
||||
"snippet": h.Snippet,
|
||||
"score": h.Score,
|
||||
"source": h.Source,
|
||||
})
|
||||
}
|
||||
if h.PublishedDate != "" {
|
||||
row["published_date"] = h.PublishedDate
|
||||
}
|
||||
results = append(results, row)
|
||||
}
|
||||
writeJSON(w, 200, map[string]any{
|
||||
"results": results,
|
||||
@@ -132,8 +167,8 @@ func (s *Server) handleSearch(w http.ResponseWriter, r *http.Request, v2 bool) {
|
||||
})
|
||||
}
|
||||
|
||||
func (s *Server) logSearch(c *auth.Caller, query string, res search.RunResult) {
|
||||
qLen := len([]rune(query))
|
||||
func (s *Server) logSearch(c *auth.Caller, queries []string, f search.Filter, res search.RunResult) {
|
||||
qLen := search.QueryRuneLen(queries)
|
||||
cache := "miss"
|
||||
if res.Cached {
|
||||
cache = "hit"
|
||||
@@ -143,5 +178,10 @@ func (s *Server) logSearch(c *auth.Caller, query string, res search.RunResult) {
|
||||
s.Log.Info("search", "caller", c.Name, "q_len", qLen, "sources", src, "n", len(res.Hits), "cache", cache)
|
||||
return
|
||||
}
|
||||
s.Log.Info("search", "caller", c.Name, "q_len", qLen, "sources", src, "n", len(res.Hits), "cache", cache, "rerank_ms", res.RerankMS)
|
||||
kinds := search.FilterKinds(f)
|
||||
attrs := []any{"caller", c.Name, "q_len", qLen, "sources", src, "n", len(res.Hits), "cache", cache, "rerank_ms", res.RerankMS, "n_queries", len(queries)}
|
||||
if kinds != "" {
|
||||
attrs = append(attrs, "filters", kinds)
|
||||
}
|
||||
s.Log.Info("search", attrs...)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,117 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/search"
|
||||
)
|
||||
|
||||
type stubSearch struct {
|
||||
hits []search.Hit
|
||||
}
|
||||
|
||||
func (s stubSearch) Name() string { return search.NameSearxNG }
|
||||
func (s stubSearch) Search(_ context.Context, _ string, _ int, _ search.Filter) search.SourceOut {
|
||||
hits := s.hits
|
||||
if hits == nil {
|
||||
hits = []search.Hit{{URL: "https://ok.example/a", Title: "t", Snippet: "s", Source: search.NameSearxNG, PublishedDate: "2026-09-01"}}
|
||||
}
|
||||
return search.SourceOut{Hits: hits}
|
||||
}
|
||||
|
||||
func TestSearchValidationTable(t *testing.T) {
|
||||
s, key, _, _, _ := searchServer(t)
|
||||
s.Search.Sources = []search.Source{stubSearch{}}
|
||||
cases := []struct {
|
||||
body map[string]any
|
||||
err string
|
||||
hint string
|
||||
}{
|
||||
{map[string]any{"query": ""}, "invalid query", ""},
|
||||
{map[string]any{"query": "a", "queries": []string{"b"}}, "invalid query", ""},
|
||||
{map[string]any{"queries": []string{"a", "b", "c", "d", "e"}}, "invalid query", ""},
|
||||
{map[string]any{"queries": []string{"a", ""}}, "invalid query", ""},
|
||||
{map[string]any{"query": "q", "include_domains": []string{"a.com OR b.com"}}, "invalid search", "include_domains"},
|
||||
{map[string]any{"query": "q", "include_domains_mode": "prefer"}, "invalid search", "include_domains_mode"},
|
||||
{map[string]any{"query": "q", "time_range": "hour"}, "invalid search", "time_range"},
|
||||
{map[string]any{"query": "q", "time_range": "day", "start_date": "2026-09-01"}, "invalid search", "time_range"},
|
||||
{map[string]any{"query": "q", "start_date": "2026-09-10", "end_date": "2026-09-01"}, "invalid search", "end_date"},
|
||||
{map[string]any{"query": "q", "topic": "sports"}, "invalid search", "topic"},
|
||||
{map[string]any{"query": "q", "min_score": -1}, "invalid search", "min_score"},
|
||||
{map[string]any{"query": "q", "min_score": 1.1}, "invalid search", "min_score"},
|
||||
{map[string]any{"query": "no quotes", "exact_match": true}, "invalid search", "exact_match"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
rec := postJSON(t, s, "/v1/search", key, tc.body)
|
||||
if rec.Code != 400 {
|
||||
t.Fatalf("%v: %d %s", tc.body, rec.Code, rec.Body.String())
|
||||
}
|
||||
var env map[string]any
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &env); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if env["error"] != tc.err {
|
||||
t.Fatalf("%v error %v want %q", tc.body, env["error"], tc.err)
|
||||
}
|
||||
h, _ := env["hint"].(string)
|
||||
if h != tc.hint {
|
||||
t.Fatalf("%v hint %q want %q", tc.body, h, tc.hint)
|
||||
}
|
||||
}
|
||||
rec := postJSON(t, s, "/v1/search", key, map[string]any{"query": "q", "include_domains": []string{"oracle.com"}})
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("oracle.com %d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
rec = postJSON(t, s, "/v1/search", key, map[string]any{"query": `find "phrase"`, "exact_match": true})
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("exact_match quoted %d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestV2SearchNoPublishedDate(t *testing.T) {
|
||||
s, key, _, _, _ := searchServer(t)
|
||||
s.Search.Sources = []search.Source{stubSearch{}}
|
||||
rec := postJSON(t, s, "/v2/search", key, map[string]any{"query": "contract", "limit": 3, "include_domains": []string{"secret.example"}})
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("%d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if strings.Contains(rec.Body.String(), "published_date") {
|
||||
t.Fatalf("v2 leaked published_date: %s", rec.Body.String())
|
||||
}
|
||||
var env map[string]any
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &env); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
data, _ := env["data"].(map[string]any)
|
||||
web, _ := data["web"].([]any)
|
||||
if len(web) < 1 {
|
||||
t.Fatalf("%v", env)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSearchFilterLogsNoSecrets(t *testing.T) {
|
||||
s, fleet, gs, h, _ := searchServer(t)
|
||||
s.Search.Sources = []search.Source{stubSearch{}}
|
||||
body := map[string]any{
|
||||
"queries": []string{"alpha secret", "beta secret", "gamma secret"},
|
||||
"include_domains": []string{"secret.example"},
|
||||
"rerank": false,
|
||||
}
|
||||
rec := postJSON(t, s, "/v1/search", gs, body)
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("gs %d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
rec = postJSON(t, s, "/v1/search", fleet, body)
|
||||
if rec.Code != 200 {
|
||||
t.Fatalf("fleet %d %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
dump := h.dump()
|
||||
for _, bad := range []string{"secret.example", "alpha secret", "beta secret", "gamma secret"} {
|
||||
if strings.Contains(dump, bad) {
|
||||
t.Fatalf("log leak %q:\n%s", bad, dump)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,232 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/auth"
|
||||
"git.wntrmute.dev/kyle/sift/internal/batch"
|
||||
"git.wntrmute.dev/kyle/sift/internal/clean"
|
||||
"git.wntrmute.dev/kyle/sift/internal/digest"
|
||||
"git.wntrmute.dev/kyle/sift/internal/fetch"
|
||||
"git.wntrmute.dev/kyle/sift/internal/fleetctx"
|
||||
"git.wntrmute.dev/kyle/sift/internal/pdfdoc"
|
||||
"git.wntrmute.dev/kyle/sift/internal/recover"
|
||||
)
|
||||
|
||||
// fetchTarget runs one target's per-target pipeline: fetch → recover → PDF render or clean →
|
||||
// question limit → digest → context → log line → envelope. It is the single body both /v1/fetch
|
||||
// and the batch worker drive, so it must not change what v1fetch returned.
|
||||
func (s *Server) fetchTarget(ctx context.Context, c *auth.Caller, t batch.Target) (batch.Result, error) {
|
||||
q := t.Options.Question
|
||||
if len([]rune(q)) > 500 {
|
||||
rns := []rune(q)
|
||||
q = string(rns[:500])
|
||||
}
|
||||
mode := t.Options.Mode
|
||||
if mode == "" {
|
||||
if q != "" {
|
||||
mode = "digest"
|
||||
} else {
|
||||
mode = "clean"
|
||||
}
|
||||
}
|
||||
if mode != "clean" && mode != "raw" && mode != "digest" {
|
||||
return batch.Result{}, &fetch.Error{Status: 400, Msg: "invalid url"}
|
||||
}
|
||||
wantCtx := q != ""
|
||||
if t.Options.Context != nil {
|
||||
wantCtx = *t.Options.Context
|
||||
}
|
||||
|
||||
t0 := time.Now()
|
||||
onlyMain := true
|
||||
res, cl, doc, ferr := s.doFetch(ctx, c, t.URL, t.Options.JQ, onlyMain)
|
||||
var snap *recover.Result
|
||||
if ferr != nil && t.Options.Recover && recoverWorthy(ferr) {
|
||||
sr, rerr := recover.Try(ctx, s.Fetch, t.URL)
|
||||
if rerr == nil && sr != nil && sr.Res != nil {
|
||||
res, cl, doc, ferr = s.cleanFetched(sr.Res, t.Options.JQ, onlyMain)
|
||||
snap = sr
|
||||
}
|
||||
}
|
||||
if ferr != nil {
|
||||
return batch.Result{}, ferr
|
||||
}
|
||||
fetchMS := time.Since(t0).Milliseconds()
|
||||
|
||||
if doc == nil && (t.Options.Pages != "" || t.Options.Grep != "") {
|
||||
return batch.Result{}, &fetch.Error{Status: 400, Msg: "pages and grep apply to PDFs only"}
|
||||
}
|
||||
|
||||
var docLen, pdfPages int
|
||||
var content string
|
||||
var pdfMap map[string]any
|
||||
if doc != nil {
|
||||
pdfPages = doc.Pages
|
||||
if mode == "raw" {
|
||||
return batch.Result{}, &fetch.Error{Status: 400, Msg: "raw mode does not apply to PDFs"}
|
||||
}
|
||||
rp, err := s.renderPDF(doc, pdfdoc.Options{Pages: t.Options.Pages, Grep: t.Options.Grep, OutCap: s.OutCap})
|
||||
if err != nil {
|
||||
if errors.Is(err, pdfdoc.ErrBadPages) || errors.Is(err, pdfdoc.ErrBadGrep) {
|
||||
return batch.Result{}, &fetch.Error{Status: 400, Msg: err.Error()}
|
||||
}
|
||||
return batch.Result{}, err
|
||||
}
|
||||
full := pdfdoc.Full(doc)
|
||||
docLen = utf8.RuneCountInString(full)
|
||||
if mode == "digest" {
|
||||
content = full
|
||||
} else {
|
||||
content = rp.Content
|
||||
}
|
||||
pdfMap = map[string]any{
|
||||
"pages": doc.Pages,
|
||||
"title": doc.Title,
|
||||
"outline": doc.Outline,
|
||||
"extracted": len(doc.Text),
|
||||
"first": rp.First,
|
||||
"last": rp.Last,
|
||||
"map": rp.Map,
|
||||
"grep_hits": rp.GrepHits,
|
||||
"note": rp.Note,
|
||||
}
|
||||
} else {
|
||||
content = cl.Markdown
|
||||
if cl.Kind == "json" {
|
||||
content = cl.HTML
|
||||
}
|
||||
if mode == "raw" {
|
||||
content = string(res.Body)
|
||||
if len(content) > s.OutCap {
|
||||
content = clean.TruncateUTF8(content, s.OutCap)
|
||||
cl.Trunc = true
|
||||
cl.Warning = "truncated to output cap"
|
||||
}
|
||||
}
|
||||
}
|
||||
warn := []string{}
|
||||
if cl.Warning != "" {
|
||||
warn = append(warn, cl.Warning)
|
||||
}
|
||||
|
||||
var digMeta map[string]any
|
||||
var digestMS int64
|
||||
if mode == "digest" {
|
||||
if doc != nil {
|
||||
n := len(digest.Split(pdfdoc.Full(doc), digest.ChunkTokens))
|
||||
limit := s.PDFMaxQuestionChunks
|
||||
if limit == 0 {
|
||||
limit = DefaultPDFMaxQuestionChunks
|
||||
}
|
||||
if n > limit {
|
||||
return batch.Result{}, &fetch.Error{Status: 422, Msg: fmt.Sprintf("document too long for question (%d chunks); use grep or pages", n)}
|
||||
}
|
||||
}
|
||||
if s.Digest == nil {
|
||||
warn = append(warn, "digest fell back to clean: unset")
|
||||
mode = "clean"
|
||||
content, warn = capDigestFallback(content, warn)
|
||||
digMeta = map[string]any{"fallback": "unset"}
|
||||
} else {
|
||||
dr := s.Digest.Run(ctx, content, q, res.FinalURL, s.callerClass(c))
|
||||
digestMS = dr.DigestMS
|
||||
if !dr.Cached {
|
||||
s.Prom.ObserveDigest(time.Duration(dr.DigestMS) * time.Millisecond)
|
||||
}
|
||||
if dr.Fallback != "" {
|
||||
warn = append(warn, "digest fell back to clean: "+dr.Fallback)
|
||||
mode = "clean"
|
||||
content, warn = capDigestFallback(content, warn)
|
||||
digMeta = map[string]any{"fallback": dr.Fallback, "model": dr.Model}
|
||||
} else {
|
||||
content = dr.Text
|
||||
digMeta = map[string]any{
|
||||
"model": dr.Model,
|
||||
"chunks_used": dr.ChunksUsed,
|
||||
"chunks_total": dr.ChunksTotal,
|
||||
"retrieval": dr.Retrieval,
|
||||
"embedded": dr.Embedded,
|
||||
"fallback": false,
|
||||
"cached": dr.Cached,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var ctxItems []fleetctx.Item
|
||||
var ctxMS int64
|
||||
if wantCtx && s.Ctx != nil {
|
||||
ct0 := time.Now()
|
||||
query := fleetctx.Query(q, cl.Title, headings(cl.Markdown))
|
||||
items, cw := s.Ctx.Search(ctx, c.Workspace, c.Name, query)
|
||||
ctxMS = time.Since(ct0).Milliseconds()
|
||||
if cw != "" {
|
||||
warn = append(warn, cw)
|
||||
}
|
||||
ctxItems = items
|
||||
}
|
||||
|
||||
qLen := len(q)
|
||||
if qLen > 80 {
|
||||
qLen = 80
|
||||
}
|
||||
fb, _ := "", false
|
||||
if digMeta != nil {
|
||||
if s, ok := digMeta["fallback"].(string); ok {
|
||||
fb = s
|
||||
}
|
||||
}
|
||||
s.Log.Info("fetch", "caller", c.Name, "host", s.logHost(c, res.FinalURL), "mode", mode,
|
||||
"q_len", qLen, "raw_tokens", len(res.Body)/4, "returned_tokens", len(content)/4,
|
||||
"pdf_pages", pdfPages, "digest_ms", digestMS, "fallback", fb)
|
||||
|
||||
src := map[string]any{
|
||||
"url": res.URL,
|
||||
"final_url": res.FinalURL,
|
||||
"status": res.Status,
|
||||
"content_type": res.ContentType,
|
||||
"fetched_at": res.FetchedAt.Format(time.RFC3339),
|
||||
"hash": res.SHA256(),
|
||||
"bytes": len(res.Body),
|
||||
"cached": len(res.Body) == 0,
|
||||
"kind": "live",
|
||||
}
|
||||
if snap != nil {
|
||||
src["kind"] = "snapshot"
|
||||
src["via"] = string(snap.Via)
|
||||
src["original_url"] = snap.OriginalURL
|
||||
if snap.SnapshotAt != "" {
|
||||
src["snapshot_at"] = snap.SnapshotAt
|
||||
} else {
|
||||
src["snapshot_at"] = nil
|
||||
warn = append(warn, "snapshot date unknown")
|
||||
}
|
||||
}
|
||||
m := map[string]any{
|
||||
"mode": mode,
|
||||
"untrusted": true,
|
||||
"source": src,
|
||||
"warnings": warn,
|
||||
"sizes": map[string]int{
|
||||
"raw_tokens": docLen / 4,
|
||||
"clean_tokens": digest.EstTokens(cl.Markdown),
|
||||
"returned_tokens": digest.EstTokens(content),
|
||||
},
|
||||
"timings": map[string]int64{"fetch_ms": fetchMS, "digest_ms": digestMS, "context_ms": ctxMS},
|
||||
}
|
||||
if pdfMap != nil {
|
||||
m["pdf"] = pdfMap
|
||||
}
|
||||
if digMeta != nil {
|
||||
m["digest"] = digMeta
|
||||
}
|
||||
if wantCtx {
|
||||
m["context"] = ctxItems
|
||||
}
|
||||
return batch.Result{Content: content, Meta: m}, nil
|
||||
}
|
||||
@@ -0,0 +1,394 @@
|
||||
// Package batch is the batch runner core: it owns job lifecycle, idempotency and
|
||||
// the per-caller / global caps. It knows nothing about HTTP, the guard or the real
|
||||
// fetch pipeline — per-target work is an injected FetchFunc, and persistence is the
|
||||
// dcache store. Jobs never resume after a restart, so the caller is held in memory,
|
||||
// never written to the store.
|
||||
package batch
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/auth"
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
)
|
||||
|
||||
type Options struct {
|
||||
Mode string `json:"mode,omitempty"`
|
||||
Question string `json:"question,omitempty"`
|
||||
Pages string `json:"pages,omitempty"`
|
||||
Grep string `json:"grep,omitempty"`
|
||||
JQ string `json:"jq,omitempty"`
|
||||
Context *bool `json:"context,omitempty"` // nil = /v1/fetch's default (on when there is a question)
|
||||
Recover bool `json:"recover,omitempty"`
|
||||
}
|
||||
|
||||
type Target struct {
|
||||
URL string `json:"url"`
|
||||
Options Options `json:"options"`
|
||||
}
|
||||
|
||||
// Result is a finished target: the /v1/fetch envelope's content, and every other envelope field
|
||||
// in Meta. It is stored JSON-encoded in the target's result column.
|
||||
type Result struct {
|
||||
Content string `json:"content"`
|
||||
Meta map[string]any `json:"meta,omitempty"`
|
||||
}
|
||||
|
||||
// TargetError is a failure with an agent-facing hint. Status is the HTTP status the same
|
||||
// failure gets from /v1/fetch; the runner ignores it (the metrics hook reads it). Any other
|
||||
// error is stored with its Error() text and no hint.
|
||||
type TargetError struct {
|
||||
Status int
|
||||
Msg, Hint string
|
||||
}
|
||||
|
||||
func (e *TargetError) Error() string { return e.Msg }
|
||||
|
||||
// FetchFunc does one target for the caller that submitted the job.
|
||||
type FetchFunc func(ctx context.Context, c *auth.Caller, t Target) (Result, error)
|
||||
|
||||
type Job struct {
|
||||
ID string
|
||||
Caller string
|
||||
Status string // queued | running | done | partial | failed | cancelled
|
||||
Error string
|
||||
TargetsTotal int
|
||||
TargetsDone int
|
||||
CreatedAt int64 // unix seconds
|
||||
}
|
||||
|
||||
type TargetStatus struct {
|
||||
Idx int
|
||||
URL string
|
||||
Status string // pending | running | done | failed
|
||||
Result *Result
|
||||
Error string
|
||||
Hint string
|
||||
}
|
||||
|
||||
// Hooks observe the runner for metrics. Any nil field is skipped.
|
||||
type Hooks struct {
|
||||
Job func(outcome string) // created, reused, done, partial, failed, cancelled
|
||||
Target func(err error, d time.Duration, resultBytes int) // once per stored target; err nil on success
|
||||
JobSeconds func(d time.Duration) // submit → terminal, once per job
|
||||
Depth func(live int) // live jobs, after every change
|
||||
Inflight func(n int) // targets being fetched, after every change
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
TTL time.Duration // 0 → 24h
|
||||
TargetsPerJob int // 0 → 32
|
||||
JobsPerCaller int // 0 → 1
|
||||
JobsTotal int // 0 → 8
|
||||
Workers int // 0 → 1; values above 2 are clamped to 2
|
||||
MaxBytes int64 // 0 → 64 MiB
|
||||
RetryWait time.Duration // 0 → 2s: the pause before a busy target is tried again
|
||||
RetryFor time.Duration // 0 → 10m: how long one target may keep coming back busy
|
||||
Hooks Hooks
|
||||
}
|
||||
|
||||
var (
|
||||
ErrNoTargets = errors.New("no targets")
|
||||
ErrTooManyTargets = errors.New("too many targets")
|
||||
ErrCallerBusy = errors.New("busy")
|
||||
ErrBusy = errors.New("jobs busy")
|
||||
ErrRetryLater = errors.New("retry later") // a FetchFunc wraps this when the target was only busy
|
||||
)
|
||||
|
||||
// terminal statuses mark a job as finished; only a queued job can be cancelled.
|
||||
func isJobTerminal(status string) bool {
|
||||
switch status {
|
||||
case "done", "partial", "failed", "cancelled":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
type Runner struct {
|
||||
db *dcache.DB
|
||||
fetch FetchFunc
|
||||
log *slog.Logger
|
||||
cfg Config
|
||||
closeOnce sync.Once
|
||||
submitMu sync.Mutex // serialises Submit: idempotency lookup, caps and create are one step
|
||||
|
||||
mu sync.Mutex // guards everything below
|
||||
cond *sync.Cond // on mu
|
||||
live []*liveJob // the rotation, in submit order
|
||||
cursor int // where the next scan starts, within [0, len(live)]; read as cursor % len(live)
|
||||
closed bool
|
||||
ctx context.Context // the FetchFunc context; Close cancels it
|
||||
stop context.CancelFunc
|
||||
wg sync.WaitGroup // the workers
|
||||
inflight int
|
||||
}
|
||||
|
||||
func New(db *dcache.DB, fetch FetchFunc, cfg Config, log *slog.Logger) *Runner {
|
||||
// Zero or negative means the default: a negative TTL or byte budget would make the sweep or
|
||||
// the eviction delete everything, and negative workers would panic the WaitGroup.
|
||||
if cfg.TTL <= 0 {
|
||||
cfg.TTL = 24 * time.Hour
|
||||
}
|
||||
if cfg.TargetsPerJob <= 0 {
|
||||
cfg.TargetsPerJob = 32
|
||||
}
|
||||
if cfg.JobsPerCaller <= 0 {
|
||||
cfg.JobsPerCaller = 1
|
||||
}
|
||||
if cfg.JobsTotal <= 0 {
|
||||
cfg.JobsTotal = 8
|
||||
}
|
||||
if cfg.Workers <= 0 {
|
||||
cfg.Workers = 1
|
||||
}
|
||||
if cfg.Workers > 2 {
|
||||
cfg.Workers = 2
|
||||
}
|
||||
if cfg.MaxBytes <= 0 {
|
||||
cfg.MaxBytes = 64 << 20
|
||||
}
|
||||
if cfg.RetryWait <= 0 {
|
||||
cfg.RetryWait = 2 * time.Second
|
||||
}
|
||||
if cfg.RetryFor <= 0 {
|
||||
cfg.RetryFor = 10 * time.Minute
|
||||
}
|
||||
ctx, stop := context.WithCancel(context.Background())
|
||||
r := &Runner{
|
||||
db: db,
|
||||
fetch: fetch,
|
||||
log: log,
|
||||
cfg: cfg,
|
||||
ctx: ctx,
|
||||
stop: stop,
|
||||
}
|
||||
r.cond = sync.NewCond(&r.mu)
|
||||
r.wg.Add(cfg.Workers)
|
||||
for i := 0; i < cfg.Workers; i++ {
|
||||
go r.worker()
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
// Submit stores the job and its targets before it returns. The order matters: an empty
|
||||
// request, then idempotency (which never touches the caps), then the caps, then the write.
|
||||
func (r *Runner) Submit(ctx context.Context, c *auth.Caller, callerClass string, targets []Target) (id string, reused bool, err error) {
|
||||
if len(targets) == 0 {
|
||||
return "", false, ErrNoTargets
|
||||
}
|
||||
|
||||
raw, err := json.Marshal(targets)
|
||||
if err != nil {
|
||||
return "", false, fmt.Errorf("batch: marshal targets: %w", err)
|
||||
}
|
||||
sum := sha256.Sum256([]byte(c.Name + "\x00" + string(raw)))
|
||||
idemKey := hex.EncodeToString(sum[:])
|
||||
|
||||
// The lookup, the caps and the create must be one step, or concurrent submits all pass the
|
||||
// checks before any of them is stored.
|
||||
r.submitMu.Lock()
|
||||
defer r.submitMu.Unlock()
|
||||
|
||||
now := time.Now().Unix()
|
||||
if existing, ok := r.db.BatchJobByIDem(c.Name, idemKey, now, r.cfg.TTL); ok {
|
||||
r.callJob("reused")
|
||||
return existing, true, nil
|
||||
}
|
||||
|
||||
if len(targets) > r.cfg.TargetsPerJob {
|
||||
return "", false, ErrTooManyTargets
|
||||
}
|
||||
if r.db.LiveBatchJobs(c.Name) >= r.cfg.JobsPerCaller {
|
||||
return "", false, ErrCallerBusy
|
||||
}
|
||||
if r.db.LiveBatchJobsTotal() >= r.cfg.JobsTotal {
|
||||
return "", false, ErrBusy
|
||||
}
|
||||
|
||||
id, err = randomID()
|
||||
if err != nil {
|
||||
return "", false, fmt.Errorf("batch: job id: %w", err)
|
||||
}
|
||||
|
||||
if err := r.db.CreateBatchJob(dcache.BatchJob{
|
||||
ID: id,
|
||||
Caller: c.Name,
|
||||
CallerClass: callerClass,
|
||||
IDemKey: idemKey,
|
||||
Status: "queued",
|
||||
TargetsTotal: len(targets),
|
||||
CreatedAt: now,
|
||||
Lease: now,
|
||||
}); err != nil {
|
||||
r.log.Warn("batch: create job", "job", id, "err", err)
|
||||
return "", false, fmt.Errorf("batch: create job: %w", err)
|
||||
}
|
||||
|
||||
btargets := make([]dcache.BatchTarget, len(targets))
|
||||
for i := range targets {
|
||||
opts, _ := json.Marshal(targets[i].Options)
|
||||
btargets[i] = dcache.BatchTarget{
|
||||
JobID: id,
|
||||
Idx: i,
|
||||
URL: targets[i].URL,
|
||||
Options: string(opts),
|
||||
Status: "pending",
|
||||
}
|
||||
}
|
||||
if err := r.db.PutBatchTargets(id, btargets); err != nil {
|
||||
r.log.Warn("batch: store targets", "job", id, "err", err)
|
||||
// Never leave a queued job that no worker will run: it would hold the caller's cap.
|
||||
_ = r.db.SetBatchJobStatus(id, "failed", "store", now)
|
||||
return "", false, fmt.Errorf("batch: store targets: %w", err)
|
||||
}
|
||||
|
||||
r.mu.Lock()
|
||||
r.live = append(r.live, &liveJob{
|
||||
id: id,
|
||||
caller: c,
|
||||
targets: targets,
|
||||
created: time.Unix(now, 0),
|
||||
})
|
||||
r.callDepth(len(r.live))
|
||||
r.cond.Broadcast()
|
||||
r.mu.Unlock()
|
||||
|
||||
r.callJob("created")
|
||||
r.db.SweepBatch(now, r.cfg.TTL)
|
||||
return id, false, nil
|
||||
}
|
||||
|
||||
// Cancel cancels a job. A job still in the rotation (queued or with a target in flight) is marked
|
||||
// cancelled at once: its in-flight target, if any, finishes and is stored without overwriting
|
||||
// cancelled, and no further target starts. A job not in the rotation (unknown or already terminal)
|
||||
// keeps task 02's behaviour.
|
||||
func (r *Runner) Cancel(id string) bool {
|
||||
r.mu.Lock()
|
||||
var j *liveJob
|
||||
for _, lj := range r.live {
|
||||
if lj.id == id {
|
||||
j = lj
|
||||
break
|
||||
}
|
||||
}
|
||||
if j != nil {
|
||||
if j.cancelled { // a concurrent DELETE got here first
|
||||
r.mu.Unlock()
|
||||
return false
|
||||
}
|
||||
j.cancelled = true
|
||||
if err := r.db.SetBatchJobStatus(id, "cancelled", "", time.Now().Unix()); err != nil {
|
||||
r.log.Warn("batch: cancel job", "job", id, "err", err)
|
||||
}
|
||||
r.callJob("cancelled")
|
||||
if !j.running {
|
||||
r.removeLocked(j)
|
||||
r.callDepth(len(r.live))
|
||||
}
|
||||
r.cond.Broadcast()
|
||||
r.mu.Unlock()
|
||||
return true
|
||||
}
|
||||
r.mu.Unlock()
|
||||
|
||||
jb, ok := r.db.BatchJob(id)
|
||||
if !ok || isJobTerminal(jb.Status) {
|
||||
return false
|
||||
}
|
||||
if err := r.db.SetBatchJobStatus(id, "cancelled", "", time.Now().Unix()); err != nil {
|
||||
r.log.Warn("batch: cancel job", "job", id, "err", err)
|
||||
return false
|
||||
}
|
||||
r.callJob("cancelled")
|
||||
return true
|
||||
}
|
||||
|
||||
// Status maps the store rows for a job and its targets; Result is decoded from the result
|
||||
// column. It writes nothing.
|
||||
func (r *Runner) Status(id string) (Job, []TargetStatus, bool) {
|
||||
jb, ok := r.db.BatchJob(id)
|
||||
if !ok {
|
||||
return Job{}, nil, false
|
||||
}
|
||||
j := Job{
|
||||
ID: jb.ID,
|
||||
Caller: jb.Caller,
|
||||
Status: jb.Status,
|
||||
Error: jb.Error,
|
||||
TargetsTotal: jb.TargetsTotal,
|
||||
TargetsDone: jb.TargetsDone,
|
||||
CreatedAt: jb.CreatedAt,
|
||||
}
|
||||
tbs, _ := r.db.BatchTargets(id)
|
||||
ts := make([]TargetStatus, 0, len(tbs))
|
||||
for _, bt := range tbs {
|
||||
var res *Result
|
||||
if bt.Result != nil {
|
||||
res = new(Result)
|
||||
if err := json.Unmarshal([]byte(*bt.Result), res); err != nil {
|
||||
r.log.Warn("batch: decode result", "job", id, "err", err)
|
||||
}
|
||||
}
|
||||
ts = append(ts, TargetStatus{
|
||||
Idx: bt.Idx,
|
||||
URL: bt.URL,
|
||||
Status: bt.Status,
|
||||
Result: res,
|
||||
Error: strOrEmpty(bt.Error),
|
||||
Hint: strOrEmpty(bt.Hint),
|
||||
})
|
||||
}
|
||||
return j, ts, true
|
||||
}
|
||||
|
||||
// Touch bumps the job's lease to now so a slow owner does not lose it.
|
||||
func (r *Runner) Touch(id string) {
|
||||
now := time.Now().Unix()
|
||||
if err := r.db.TouchBatchJob(id, now); err != nil {
|
||||
r.log.Warn("batch: touch job", "job", id, "err", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Close stops the workers, cancels their context, and waits for an in-flight fetch to return.
|
||||
// Safe to call more than once. Jobs still live stay queued or running in the store, for
|
||||
// MarkInFlightFailed("restart") on the next start.
|
||||
func (r *Runner) Close() {
|
||||
r.closeOnce.Do(func() {
|
||||
r.mu.Lock()
|
||||
r.closed = true
|
||||
r.stop()
|
||||
r.cond.Broadcast()
|
||||
r.mu.Unlock()
|
||||
r.wg.Wait()
|
||||
})
|
||||
}
|
||||
|
||||
func (r *Runner) callJob(outcome string) {
|
||||
if r.cfg.Hooks.Job != nil {
|
||||
r.cfg.Hooks.Job(outcome)
|
||||
}
|
||||
}
|
||||
|
||||
func randomID() (string, error) {
|
||||
var b [16]byte
|
||||
if _, err := rand.Read(b[:]); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return hex.EncodeToString(b[:]), nil
|
||||
}
|
||||
|
||||
func strOrEmpty(s *string) string {
|
||||
if s == nil {
|
||||
return ""
|
||||
}
|
||||
return *s
|
||||
}
|
||||
@@ -0,0 +1,177 @@
|
||||
package batch
|
||||
|
||||
// P6 owner review: regression tests for the runner (findings from the independent review).
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/auth"
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
)
|
||||
|
||||
// A job cancelled while its in-flight target comes back busy leaves the rotation.
|
||||
func TestCancelWhileBusyLeavesRotation(t *testing.T) {
|
||||
started := make(chan struct{}, 1)
|
||||
release := make(chan struct{})
|
||||
var depth atomic.Int64
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
select {
|
||||
case started <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
<-release
|
||||
return Result{}, fmt.Errorf("busy: %w", ErrRetryLater)
|
||||
}
|
||||
r := newRunner(t, fetch, Config{RetryWait: time.Millisecond, Hooks: Hooks{Depth: func(n int) { depth.Store(int64(n)) }}})
|
||||
id, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
<-started
|
||||
if !r.Cancel(id) {
|
||||
t.Fatal("cancel returned false")
|
||||
}
|
||||
close(release)
|
||||
deadline := time.Now().Add(2 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
r.mu.Lock()
|
||||
n := len(r.live)
|
||||
r.mu.Unlock()
|
||||
if n == 0 && depth.Load() == 0 {
|
||||
return
|
||||
}
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
t.Fatalf("cancelled job still in the rotation (depth %d)", depth.Load())
|
||||
}
|
||||
|
||||
// Close stops the workers: no further target is fetched, and the job is left live in the
|
||||
// store for MarkInFlightFailed("restart") to reap, not failed target by target.
|
||||
func TestCloseStopsDraining(t *testing.T) {
|
||||
started := make(chan struct{}, 1)
|
||||
var calls atomic.Int64
|
||||
fetch := func(ctx context.Context, _ *auth.Caller, _ Target) (Result, error) {
|
||||
calls.Add(1)
|
||||
select {
|
||||
case started <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
<-ctx.Done()
|
||||
return Result{}, ctx.Err()
|
||||
}
|
||||
r := newRunner(t, fetch, Config{})
|
||||
id, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b", "c", "d"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
<-started
|
||||
r.Close()
|
||||
if n := calls.Load(); n != 1 {
|
||||
t.Fatalf("%d fetches after Close, want only the one in flight", n)
|
||||
}
|
||||
j, ts, _ := r.Status(id)
|
||||
if j.Status != "running" && j.Status != "queued" {
|
||||
t.Fatalf("job %q after Close, want it left live for the restart reaper", j.Status)
|
||||
}
|
||||
for _, x := range ts {
|
||||
if x.Status != "pending" {
|
||||
t.Fatalf("target %d is %q after Close, want pending", x.Idx, x.Status)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Concurrent submits respect the caps and idempotency: one job per caller, and identical
|
||||
// lists all get the same job.
|
||||
func TestConcurrentSubmits(t *testing.T) {
|
||||
release := make(chan struct{})
|
||||
defer close(release)
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
<-release
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{})
|
||||
var wg sync.WaitGroup
|
||||
var mu sync.Mutex
|
||||
ids := map[string]bool{}
|
||||
for i := 0; i < 20; i++ {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
id, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls(fmt.Sprintf("u%d", i)))
|
||||
if err == nil {
|
||||
mu.Lock()
|
||||
ids[id] = true
|
||||
mu.Unlock()
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
if len(ids) != 1 {
|
||||
t.Fatalf("20 concurrent submits by one caller (cap 1) created %d jobs", len(ids))
|
||||
}
|
||||
same := map[string]bool{}
|
||||
for i := 0; i < 20; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
id, _, err := r.Submit(context.Background(), bob, dcache.ClassFleet, urls("same"))
|
||||
if err == nil {
|
||||
mu.Lock()
|
||||
same[id] = true
|
||||
mu.Unlock()
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
if len(same) != 1 {
|
||||
t.Fatalf("20 concurrent identical submits created %d jobs, want 1", len(same))
|
||||
}
|
||||
}
|
||||
|
||||
// A second cancel of a job with a target in flight returns false and counts once.
|
||||
func TestDoubleCancel(t *testing.T) {
|
||||
started := make(chan struct{}, 1)
|
||||
release := make(chan struct{})
|
||||
var cancels atomic.Int64
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
select {
|
||||
case started <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
<-release
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{Hooks: Hooks{Job: func(o string) {
|
||||
if o == "cancelled" {
|
||||
cancels.Add(1)
|
||||
}
|
||||
}}})
|
||||
defer close(release)
|
||||
id, _, _ := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
<-started
|
||||
if !r.Cancel(id) {
|
||||
t.Fatal("first cancel returned false")
|
||||
}
|
||||
if r.Cancel(id) {
|
||||
t.Fatal("second cancel returned true")
|
||||
}
|
||||
if n := cancels.Load(); n != 1 {
|
||||
t.Fatalf("cancelled counted %d times", n)
|
||||
}
|
||||
}
|
||||
|
||||
// Negative or zero durations and sizes fall back to the defaults.
|
||||
func TestConfigDefaults(t *testing.T) {
|
||||
r := newRunner(t, func(context.Context, *auth.Caller, Target) (Result, error) { return Result{}, nil },
|
||||
Config{TTL: -time.Hour, MaxBytes: -1, RetryWait: -1, RetryFor: -1, TargetsPerJob: -1, JobsPerCaller: -1, JobsTotal: -1, Workers: -1})
|
||||
c := r.cfg
|
||||
if c.TTL != 24*time.Hour || c.MaxBytes != 64<<20 || c.RetryWait != 2*time.Second || c.RetryFor != 10*time.Minute ||
|
||||
c.TargetsPerJob != 32 || c.JobsPerCaller != 1 || c.JobsTotal != 8 || c.Workers != 1 {
|
||||
t.Fatalf("config %+v", c)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,182 @@
|
||||
package batch
|
||||
|
||||
// P6 task 02: the runner's core — types, Submit (idempotency, caps, storage), Status, Touch and
|
||||
// Cancel — with no workers yet: every job stays queued. The helpers here are shared with task
|
||||
// 03's runner_test.go.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"log/slog"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/auth"
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
)
|
||||
|
||||
var (
|
||||
alice = &auth.Caller{Name: "alice", Workspace: "hermes-fleet"}
|
||||
bob = &auth.Caller{Name: "bob", Workspace: "hermes-fleet"}
|
||||
carol = &auth.Caller{Name: "carol", Workspace: "hermes-fleet"}
|
||||
)
|
||||
|
||||
func newRunner(t *testing.T, fetch FetchFunc, cfg Config) *Runner {
|
||||
t.Helper()
|
||||
db, err := dcache.Open("", 1<<20, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
r := New(db, fetch, cfg, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
t.Cleanup(r.Close) // runs before db.Close (cleanups are LIFO)
|
||||
return r
|
||||
}
|
||||
|
||||
func waitStatus(t *testing.T, r *Runner, id, want string) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
j, _, ok := r.Status(id)
|
||||
if ok && j.Status == want {
|
||||
return
|
||||
}
|
||||
time.Sleep(2 * time.Millisecond)
|
||||
}
|
||||
j, _, _ := r.Status(id)
|
||||
t.Fatalf("job %s never reached %q (now %q)", id, want, j.Status)
|
||||
}
|
||||
|
||||
func urls(ss ...string) []Target {
|
||||
var out []Target
|
||||
for _, s := range ss {
|
||||
out = append(out, Target{URL: "https://" + s + ".example/"})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Submit stores the job queued with its targets pending, before it returns.
|
||||
func TestSubmitStoresQueued(t *testing.T) {
|
||||
r := newRunner(t, func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
return Result{Content: "ok"}, nil
|
||||
}, Config{Workers: 1})
|
||||
id, reused, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
if err != nil || reused || len(id) != 32 {
|
||||
t.Fatalf("submit id=%q reused=%v err=%v, want a 32-hex-char id", id, reused, err)
|
||||
}
|
||||
j, ts, ok := r.Status(id)
|
||||
if !ok || j.ID != id || j.Caller != "alice" || j.TargetsTotal != 2 || j.CreatedAt == 0 {
|
||||
t.Fatalf("job %+v ok=%v", j, ok)
|
||||
}
|
||||
if j.Status != "queued" && j.Status != "running" && j.Status != "done" {
|
||||
t.Fatalf("status %q", j.Status)
|
||||
}
|
||||
if len(ts) != 2 || ts[0].Idx != 0 || ts[0].URL != "https://a.example/" || ts[1].URL != "https://b.example/" {
|
||||
t.Fatalf("targets %+v", ts)
|
||||
}
|
||||
}
|
||||
|
||||
// Cancelling a job no worker has reached: it is cancelled at once, its targets stay pending,
|
||||
// and the same list then starts a new job.
|
||||
func TestCancelQueued(t *testing.T) {
|
||||
release := make(chan struct{})
|
||||
defer close(release)
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
<-release
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{Workers: 1})
|
||||
busy, _, err := r.Submit(context.Background(), bob, dcache.ClassFleet, urls("x")) // occupies the only worker once there is one
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_ = busy
|
||||
id, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !r.Cancel(id) {
|
||||
t.Fatal("cancel returned false")
|
||||
}
|
||||
j, ts, _ := r.Status(id)
|
||||
if j.Status != "cancelled" || ts[0].Status != "pending" || ts[1].Status != "pending" {
|
||||
t.Fatalf("after cancel: %+v %+v", j, ts)
|
||||
}
|
||||
if r.Cancel(id) || r.Cancel("no-such-job") {
|
||||
t.Fatal("cancel of a finished or unknown job must return false")
|
||||
}
|
||||
id2, reused, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
if err != nil || reused || id2 == id {
|
||||
t.Fatalf("resubmit after cancel id=%q reused=%v err=%v", id2, reused, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Touch keeps a job alive; Status never does.
|
||||
func TestTouchAndStatusLease(t *testing.T) {
|
||||
db, err := dcache.Open("", 1<<20, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
release := make(chan struct{})
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
<-release
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := New(db, fetch, Config{Workers: 1}, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
t.Cleanup(r.Close)
|
||||
defer close(release) // runs before the cleanups
|
||||
id, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond) // let a worker (task 03) take the target; it then blocks
|
||||
before, _ := db.BatchJob(id)
|
||||
time.Sleep(1100 * time.Millisecond) // the lease is in whole seconds
|
||||
r.Status(id)
|
||||
if j, _ := db.BatchJob(id); j.Lease != before.Lease {
|
||||
t.Fatalf("Status moved the lease %d → %d; a read must not write", before.Lease, j.Lease)
|
||||
}
|
||||
r.Touch(id)
|
||||
if j, _ := db.BatchJob(id); j.Lease <= before.Lease {
|
||||
t.Fatalf("Touch did not move the lease: %d → %d", before.Lease, j.Lease)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCaps(t *testing.T) {
|
||||
release := make(chan struct{})
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
<-release
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{TargetsPerJob: 2, JobsTotal: 2})
|
||||
defer close(release)
|
||||
if _, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b", "c")); !errors.Is(err, ErrTooManyTargets) {
|
||||
t.Fatalf("3 targets over a cap of 2: want ErrTooManyTargets, got %v", err)
|
||||
}
|
||||
if _, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, nil); !errors.Is(err, ErrNoTargets) {
|
||||
t.Fatalf("no targets: want ErrNoTargets, got %v", err)
|
||||
}
|
||||
if _, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// alice's job is live (queued or running): a second, different list is refused.
|
||||
if _, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("z")); !errors.Is(err, ErrCallerBusy) {
|
||||
t.Fatalf("want ErrCallerBusy, got %v", err)
|
||||
}
|
||||
if _, _, err := r.Submit(context.Background(), bob, dcache.ClassFleet, urls("b")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Two live jobs in total: a third caller is refused.
|
||||
if _, _, err := r.Submit(context.Background(), carol, dcache.ClassFleet, urls("c")); !errors.Is(err, ErrBusy) {
|
||||
t.Fatalf("want ErrBusy, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStatusUnknown(t *testing.T) {
|
||||
r := newRunner(t, func(context.Context, *auth.Caller, Target) (Result, error) { return Result{}, nil }, Config{})
|
||||
if _, _, ok := r.Status("no-such-job"); ok {
|
||||
t.Fatal("unknown id must be ok=false")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,373 @@
|
||||
package batch
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/auth"
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
)
|
||||
|
||||
// Targets run one at a time, in order, and each fetch is given the submitting caller.
|
||||
func TestSerialOrderAndCaller(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
var calls, names []string
|
||||
inflight, peak := 0, 0
|
||||
fetch := func(_ context.Context, c *auth.Caller, tgt Target) (Result, error) {
|
||||
mu.Lock()
|
||||
inflight++
|
||||
if inflight > peak {
|
||||
peak = inflight
|
||||
}
|
||||
calls = append(calls, tgt.URL)
|
||||
names = append(names, c.Name)
|
||||
mu.Unlock()
|
||||
time.Sleep(time.Millisecond)
|
||||
mu.Lock()
|
||||
inflight--
|
||||
mu.Unlock()
|
||||
return Result{Content: "ok " + tgt.URL}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{})
|
||||
id, reused, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b", "c"))
|
||||
if err != nil || reused {
|
||||
t.Fatalf("submit err=%v reused=%v", err, reused)
|
||||
}
|
||||
waitStatus(t, r, id, "done")
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if strings.Join(calls, " ") != "https://a.example/ https://b.example/ https://c.example/" {
|
||||
t.Fatalf("order %v", calls)
|
||||
}
|
||||
if peak != 1 {
|
||||
t.Fatalf("peak inflight %d, want 1 (serial)", peak)
|
||||
}
|
||||
for _, n := range names {
|
||||
if n != "alice" {
|
||||
t.Fatalf("fetch saw caller %q, want alice", n)
|
||||
}
|
||||
}
|
||||
j, _, _ := r.Status(id)
|
||||
if j.Caller != "alice" || j.TargetsTotal != 3 || j.TargetsDone != 3 || j.CreatedAt == 0 {
|
||||
t.Fatalf("job %+v", j)
|
||||
}
|
||||
}
|
||||
|
||||
// With one worker, two live jobs take turns target by target.
|
||||
func TestRoundRobin(t *testing.T) {
|
||||
first := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
var once sync.Once
|
||||
var mu sync.Mutex
|
||||
var calls []string
|
||||
fetch := func(_ context.Context, _ *auth.Caller, tgt Target) (Result, error) {
|
||||
mu.Lock()
|
||||
calls = append(calls, strings.TrimSuffix(strings.TrimPrefix(tgt.URL, "https://"), ".example/"))
|
||||
mu.Unlock()
|
||||
once.Do(func() { close(first); <-release })
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{Workers: 1})
|
||||
a, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a1", "a2", "a3"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
<-first // a1 is in flight
|
||||
b, _, err := r.Submit(context.Background(), bob, dcache.ClassFleet, urls("b1", "b2", "b3"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
close(release)
|
||||
waitStatus(t, r, a, "done")
|
||||
waitStatus(t, r, b, "done")
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if got := strings.Join(calls, " "); got != "a1 b1 a2 b2 a3 b3" {
|
||||
t.Fatalf("order %q, want round-robin a1 b1 a2 b2 a3 b3", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFailureIsolation(t *testing.T) {
|
||||
fetch := func(_ context.Context, _ *auth.Caller, tgt Target) (Result, error) {
|
||||
if tgt.URL == "https://b.example/" {
|
||||
return Result{}, &TargetError{Msg: "upstream 500", Hint: "browser"}
|
||||
}
|
||||
return Result{Content: "ok " + tgt.URL}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{})
|
||||
id, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b", "c"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
waitStatus(t, r, id, "partial")
|
||||
j, ts, ok := r.Status(id)
|
||||
if !ok || j.TargetsDone != 3 {
|
||||
t.Fatalf("job %+v ok=%v", j, ok)
|
||||
}
|
||||
if ts[0].Status != "done" || ts[1].Status != "failed" || ts[2].Status != "done" {
|
||||
t.Fatalf("statuses %+v", ts)
|
||||
}
|
||||
if ts[1].Error != "upstream 500" || ts[1].Hint != "browser" {
|
||||
t.Fatalf("failed target error=%q hint=%q", ts[1].Error, ts[1].Hint)
|
||||
}
|
||||
if ts[0].Result == nil || ts[0].Result.Content != "ok https://a.example/" {
|
||||
t.Fatalf("result %+v", ts[0].Result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAllFailedIsFailed(t *testing.T) {
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) { return Result{}, errors.New("nope") }
|
||||
r := newRunner(t, fetch, Config{})
|
||||
id, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
waitStatus(t, r, id, "failed")
|
||||
_, ts, _ := r.Status(id)
|
||||
if ts[0].Error != "nope" || ts[0].Hint != "" {
|
||||
t.Fatalf("plain error stored as error=%q hint=%q", ts[0].Error, ts[0].Hint)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIdempotentResubmit(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
n := 0
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
mu.Lock()
|
||||
n++
|
||||
mu.Unlock()
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{})
|
||||
id1, reused1, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
if err != nil || reused1 {
|
||||
t.Fatalf("first submit err=%v reused=%v", err, reused1)
|
||||
}
|
||||
id2, reused2, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
if err != nil || !reused2 || id2 != id1 {
|
||||
t.Fatalf("resubmit id=%q reused=%v err=%v", id2, reused2, err)
|
||||
}
|
||||
waitStatus(t, r, id1, "done")
|
||||
// Still reused once done (within the lease), and never re-fetched.
|
||||
if id3, reused3, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b")); err != nil || !reused3 || id3 != id1 {
|
||||
t.Fatalf("resubmit after done id=%q reused=%v err=%v", id3, reused3, err)
|
||||
}
|
||||
// Another caller with the same list gets its own job.
|
||||
if id4, reused4, err := r.Submit(context.Background(), bob, dcache.ClassFleet, urls("a", "b")); err != nil || reused4 || id4 == id1 {
|
||||
t.Fatalf("other caller id=%q reused=%v err=%v", id4, reused4, err)
|
||||
}
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if n != 4 {
|
||||
t.Fatalf("fetched %d times, want 4 (2 for alice's one job, 2 for bob's)", n)
|
||||
}
|
||||
}
|
||||
|
||||
// A cancelled job is not handed back to a resubmit: the same list starts a new job.
|
||||
func TestResubmitAfterCancelIsNew(t *testing.T) {
|
||||
release := make(chan struct{})
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
<-release
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{})
|
||||
id1, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
waitStatus(t, r, id1, "running")
|
||||
if !r.Cancel(id1) {
|
||||
t.Fatal("cancel returned false")
|
||||
}
|
||||
close(release)
|
||||
waitStatus(t, r, id1, "cancelled")
|
||||
id2, reused, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
if err != nil || reused || id2 == id1 {
|
||||
t.Fatalf("resubmit after cancel id=%q reused=%v err=%v", id2, reused, err)
|
||||
}
|
||||
}
|
||||
|
||||
// Submit returns with the job stored "queued"; it becomes "running" when a worker takes it.
|
||||
func TestQueuedThenRunning(t *testing.T) {
|
||||
release := make(chan struct{})
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
<-release
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{Workers: 1})
|
||||
defer close(release)
|
||||
a, _, _ := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a"))
|
||||
waitStatus(t, r, a, "running")
|
||||
b, _, err := r.Submit(context.Background(), bob, dcache.ClassFleet, urls("b"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// The only worker is busy with alice's target, so bob's job waits in the queue.
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
if j, _, _ := r.Status(b); j.Status != "queued" {
|
||||
t.Fatalf("bob's job is %q while the only worker is busy, want queued", j.Status)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCancel(t *testing.T) {
|
||||
started := make(chan struct{}, 1)
|
||||
release := make(chan struct{})
|
||||
var mu sync.Mutex
|
||||
var calls []string
|
||||
fetch := func(_ context.Context, _ *auth.Caller, tgt Target) (Result, error) {
|
||||
mu.Lock()
|
||||
calls = append(calls, tgt.URL)
|
||||
mu.Unlock()
|
||||
select {
|
||||
case started <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
<-release
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{})
|
||||
id, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b", "c"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
<-started // the first target is in flight
|
||||
if !r.Cancel(id) {
|
||||
t.Fatal("cancel returned false")
|
||||
}
|
||||
close(release) // let the in-flight target finish
|
||||
waitStatus(t, r, id, "cancelled")
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
mu.Lock()
|
||||
n := len(calls)
|
||||
mu.Unlock()
|
||||
if n != 1 {
|
||||
t.Fatalf("cancelled after %d fetches, want 1", n)
|
||||
}
|
||||
j, ts, ok := r.Status(id)
|
||||
if !ok || j.Status != "cancelled" {
|
||||
t.Fatalf("job %+v: the in-flight target's write must not overwrite cancelled", j)
|
||||
}
|
||||
if ts[0].Status != "done" {
|
||||
t.Fatalf("in-flight target should be stored done: %+v", ts[0])
|
||||
}
|
||||
if ts[1].Status != "pending" || ts[2].Status != "pending" {
|
||||
t.Fatalf("later targets must stay pending: %+v", ts[1:])
|
||||
}
|
||||
if r.Cancel(id) {
|
||||
t.Fatal("cancelling a terminal job must return false")
|
||||
}
|
||||
if r.Cancel("no-such-job") {
|
||||
t.Fatal("cancelling an unknown job must return false")
|
||||
}
|
||||
}
|
||||
|
||||
// A target that reports ErrRetryLater (the fetch gate was busy) is retried, not failed.
|
||||
func TestBusyIsRetried(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
n := 0
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
n++
|
||||
if n <= 2 {
|
||||
return Result{}, fmt.Errorf("busy: %w", ErrRetryLater)
|
||||
}
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{RetryWait: time.Millisecond})
|
||||
id, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
waitStatus(t, r, id, "done")
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if n != 3 {
|
||||
t.Fatalf("fetched %d times, want 3 (two busy, then ok)", n)
|
||||
}
|
||||
}
|
||||
|
||||
// Retrying is bounded: a target busy for longer than RetryFor fails with a busy error.
|
||||
func TestBusyGivesUp(t *testing.T) {
|
||||
fetch := func(context.Context, *auth.Caller, Target) (Result, error) {
|
||||
return Result{}, fmt.Errorf("busy: %w", ErrRetryLater)
|
||||
}
|
||||
r := newRunner(t, fetch, Config{RetryWait: time.Millisecond, RetryFor: 30 * time.Millisecond})
|
||||
id, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
waitStatus(t, r, id, "failed")
|
||||
_, ts, _ := r.Status(id)
|
||||
if ts[0].Status != "failed" || !strings.Contains(ts[0].Error, "busy") {
|
||||
t.Fatalf("target %+v, want failed with a busy error", ts[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestPanicRecovered(t *testing.T) {
|
||||
fetch := func(_ context.Context, _ *auth.Caller, tgt Target) (Result, error) {
|
||||
if tgt.URL == "https://b.example/" {
|
||||
panic("boom")
|
||||
}
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{})
|
||||
id, _, err := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b", "c"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
waitStatus(t, r, id, "partial")
|
||||
_, ts, _ := r.Status(id)
|
||||
if ts[0].Status != "done" || ts[1].Status != "failed" || ts[2].Status != "done" {
|
||||
t.Fatalf("statuses %+v", ts)
|
||||
}
|
||||
if ts[1].Error == "" || strings.Contains(ts[1].Error, "boom") {
|
||||
t.Fatalf("panic target error %q: must be set, and must not echo the panic value", ts[1].Error)
|
||||
}
|
||||
}
|
||||
|
||||
// The hooks see every job outcome once and every stored target once.
|
||||
func TestHooks(t *testing.T) {
|
||||
var mu sync.Mutex
|
||||
jobs := map[string]int{}
|
||||
var targets, jobSecs int
|
||||
var bytes int
|
||||
h := Hooks{
|
||||
Job: func(o string) { mu.Lock(); jobs[o]++; mu.Unlock() },
|
||||
Target: func(err error, _ time.Duration, n int) {
|
||||
mu.Lock()
|
||||
targets++
|
||||
bytes += n
|
||||
mu.Unlock()
|
||||
},
|
||||
JobSeconds: func(time.Duration) { mu.Lock(); jobSecs++; mu.Unlock() },
|
||||
}
|
||||
fetch := func(_ context.Context, _ *auth.Caller, tgt Target) (Result, error) {
|
||||
if tgt.URL == "https://b.example/" {
|
||||
return Result{}, errors.New("nope")
|
||||
}
|
||||
return Result{Content: "ok"}, nil
|
||||
}
|
||||
r := newRunner(t, fetch, Config{Hooks: h})
|
||||
id, _, _ := r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b"))
|
||||
waitStatus(t, r, id, "partial")
|
||||
_, _, _ = r.Submit(context.Background(), alice, dcache.ClassFleet, urls("a", "b")) // reused
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if jobs["created"] != 1 || jobs["partial"] != 1 || jobs["reused"] != 1 || len(jobs) != 3 {
|
||||
t.Fatalf("job outcomes %v", jobs)
|
||||
}
|
||||
if targets != 2 || jobSecs != 1 {
|
||||
t.Fatalf("targets=%d jobSeconds=%d, want 2 and 1", targets, jobSecs)
|
||||
}
|
||||
if bytes <= 0 {
|
||||
t.Fatalf("stored result bytes %d, want > 0 for the done target", bytes)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,286 @@
|
||||
package batch
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/auth"
|
||||
)
|
||||
|
||||
// ErrInternal is stored for a target whose fetch panicked; the panic value is never echoed.
|
||||
var ErrInternal = errors.New("internal error")
|
||||
|
||||
// liveJob is a submitted job that the workers are draining. It lives in Runner.live until the
|
||||
// last target is stored or the job is cancelled.
|
||||
type liveJob struct {
|
||||
id string
|
||||
caller *auth.Caller
|
||||
targets []Target // as submitted
|
||||
created time.Time
|
||||
next int // index of the next target to fetch
|
||||
running bool // a target of this job is in flight
|
||||
started bool // set to "running" in the store already
|
||||
cancelled bool
|
||||
retryAt time.Time // zero, or when the busy target may be tried again
|
||||
busySince time.Time // zero, or when the current target first came back busy
|
||||
done, failed int
|
||||
}
|
||||
|
||||
// worker drains the rotation: it takes the next live job with a pending target and no target in
|
||||
// flight, fetches and stores that one target, and moves on. Each job is serial — its target i+1
|
||||
// never starts before i is stored.
|
||||
func (r *Runner) worker() {
|
||||
defer r.wg.Done()
|
||||
|
||||
for {
|
||||
r.mu.Lock()
|
||||
if r.closed { // stop taking targets at once; live jobs are reaped on the next start
|
||||
r.mu.Unlock()
|
||||
return
|
||||
}
|
||||
now := time.Now()
|
||||
j := r.pick(now)
|
||||
if j == nil {
|
||||
r.scheduleRetryLocked(now)
|
||||
r.cond.Wait() // re-locks on wake; release before the next top-of-loop Lock
|
||||
r.mu.Unlock()
|
||||
continue
|
||||
}
|
||||
|
||||
// Take the target: mark the job running and bump the in-flight count under the lock,
|
||||
// then fetch outside it.
|
||||
j.running = true
|
||||
r.inflight++
|
||||
r.callInflight(r.inflight)
|
||||
if !j.started {
|
||||
if err := r.db.SetBatchJobStatus(j.id, "running", "", time.Now().Unix()); err != nil {
|
||||
r.log.Warn("batch: set running", "job", j.id, "err", err)
|
||||
}
|
||||
j.started = true
|
||||
}
|
||||
idx := j.next
|
||||
target := j.targets[idx]
|
||||
startedAt := time.Now()
|
||||
r.mu.Unlock()
|
||||
|
||||
res, err := r.runTarget(r.ctx, j, idx, target)
|
||||
|
||||
r.mu.Lock()
|
||||
now = time.Now() // after the fetch: a busy answer can take seconds to come back
|
||||
j.running = false
|
||||
r.inflight--
|
||||
r.callInflight(r.inflight)
|
||||
|
||||
// Shutting down: a failure now is most likely the cancelled context. Store nothing; the
|
||||
// job stays live in the store for the restart reaper.
|
||||
if r.closed && err != nil {
|
||||
r.cond.Broadcast()
|
||||
r.mu.Unlock()
|
||||
continue
|
||||
}
|
||||
|
||||
// A cancelled job whose in-flight target was only busy is finished: nothing more runs.
|
||||
if errors.Is(err, ErrRetryLater) && j.cancelled {
|
||||
r.removeLocked(j)
|
||||
r.callDepth(len(r.live))
|
||||
r.cond.Broadcast()
|
||||
r.mu.Unlock()
|
||||
continue
|
||||
}
|
||||
|
||||
// A target that was only busy stays pending and is retried after RetryWait, unless it
|
||||
// has been busy for longer than RetryFor, in which case it is stored failed below.
|
||||
if errors.Is(err, ErrRetryLater) {
|
||||
if j.busySince.IsZero() {
|
||||
j.busySince = now
|
||||
}
|
||||
if now.Sub(j.busySince) <= r.cfg.RetryFor {
|
||||
j.retryAt = now.Add(r.cfg.RetryWait)
|
||||
r.cond.Broadcast()
|
||||
r.mu.Unlock()
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
resultBytes := r.storeTarget(j, idx, res, err)
|
||||
j.next++
|
||||
j.busySince, j.retryAt = time.Time{}, time.Time{}
|
||||
if err == nil {
|
||||
j.done++
|
||||
} else {
|
||||
j.failed++
|
||||
}
|
||||
r.callTarget(err, time.Since(startedAt), resultBytes)
|
||||
r.db.EvictBatchOldest(r.cfg.MaxBytes)
|
||||
|
||||
if j.cancelled || j.next >= len(j.targets) {
|
||||
if !j.cancelled {
|
||||
r.setTerminal(j)
|
||||
}
|
||||
r.removeLocked(j)
|
||||
r.callDepth(len(r.live))
|
||||
}
|
||||
r.cond.Broadcast()
|
||||
r.mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
// pick returns the first live job, scanning from cursor and wrapping once, that has a pending
|
||||
// target that is not running, not cancelled, and ready (now past any retry). It advances cursor to
|
||||
// the index after the job, so the rotation is round-robin across jobs in submit order.
|
||||
func (r *Runner) pick(now time.Time) *liveJob {
|
||||
n := len(r.live)
|
||||
if n == 0 {
|
||||
return nil
|
||||
}
|
||||
for i := 0; i < n; i++ {
|
||||
j := r.live[(r.cursor+i)%n]
|
||||
if j.running || j.cancelled || j.next >= len(j.targets) {
|
||||
continue
|
||||
}
|
||||
if !now.Before(j.retryAt) { // now >= retryAt; a zero retryAt is always ready
|
||||
// The job after this one; n (one past the end) means "whatever is appended next,
|
||||
// else wrap to 0" — reads take cursor % n.
|
||||
r.cursor = (r.cursor+i)%n + 1
|
||||
return j
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// scheduleRetryLocked arms a timer that wakes the workers at the earliest future retryAt, so they
|
||||
// do not sleep past a busy target that is ready to be tried again.
|
||||
func (r *Runner) scheduleRetryLocked(now time.Time) {
|
||||
var earliest time.Time
|
||||
for _, j := range r.live {
|
||||
if j.cancelled || j.next >= len(j.targets) || !j.retryAt.After(now) {
|
||||
continue
|
||||
}
|
||||
if earliest.IsZero() || j.retryAt.Before(earliest) {
|
||||
earliest = j.retryAt
|
||||
}
|
||||
}
|
||||
if !earliest.IsZero() {
|
||||
d := time.Until(earliest)
|
||||
if d < 0 {
|
||||
d = 0
|
||||
}
|
||||
time.AfterFunc(d, r.cond.Broadcast)
|
||||
}
|
||||
}
|
||||
|
||||
// removeLocked drops j from live and fixes cursor so the rotation still points at the next job.
|
||||
func (r *Runner) removeLocked(j *liveJob) {
|
||||
idx := -1
|
||||
for i := range r.live {
|
||||
if r.live[i] == j {
|
||||
idx = i
|
||||
break
|
||||
}
|
||||
}
|
||||
if idx < 0 {
|
||||
return
|
||||
}
|
||||
r.live = append(r.live[:idx], r.live[idx+1:]...)
|
||||
if r.cursor > idx {
|
||||
r.cursor--
|
||||
}
|
||||
if r.cursor > len(r.live) {
|
||||
r.cursor = len(r.live)
|
||||
}
|
||||
}
|
||||
|
||||
// setTerminal writes the job's finished status once: all done → done, all failed → failed, a mix
|
||||
// → partial.
|
||||
func (r *Runner) setTerminal(j *liveJob) {
|
||||
var status string
|
||||
switch {
|
||||
case j.failed == 0:
|
||||
status = "done"
|
||||
case j.done == 0:
|
||||
status = "failed"
|
||||
default:
|
||||
status = "partial"
|
||||
}
|
||||
if err := r.db.SetBatchJobStatus(j.id, status, "", time.Now().Unix()); err != nil {
|
||||
r.log.Warn("batch: set terminal status", "job", j.id, "err", err)
|
||||
}
|
||||
r.callJob(status)
|
||||
r.callJobSeconds(time.Since(j.created))
|
||||
}
|
||||
|
||||
// runTarget fetches one target, recovering a panic into ErrInternal (never echoing the value).
|
||||
func (r *Runner) runTarget(ctx context.Context, j *liveJob, idx int, t Target) (res Result, err error) {
|
||||
defer func() {
|
||||
if rec := recover(); rec != nil {
|
||||
r.log.Warn("batch: recover fetch panic", "job", j.id, "idx", idx)
|
||||
res = Result{}
|
||||
err = ErrInternal
|
||||
}
|
||||
}()
|
||||
return r.fetch(ctx, j.caller, t)
|
||||
}
|
||||
|
||||
// storeTarget stores one target: success → done with the result JSON, error → failed with the
|
||||
// TargetError's Msg/Hint or the error's text. It returns the number of result bytes stored.
|
||||
func (r *Runner) storeTarget(j *liveJob, idx int, res Result, err error) int {
|
||||
now := time.Now().Unix()
|
||||
status := "failed"
|
||||
var resultStr, errMsg, hint *string
|
||||
if err == nil {
|
||||
status = "done"
|
||||
b, merr := json.Marshal(res)
|
||||
if merr != nil {
|
||||
r.log.Warn("batch: marshal result", "job", j.id, "idx", idx, "err", merr)
|
||||
} else {
|
||||
s := string(b)
|
||||
resultStr = &s
|
||||
}
|
||||
} else {
|
||||
var terr *TargetError
|
||||
if errors.As(err, &terr) {
|
||||
m := terr.Msg
|
||||
errMsg = &m
|
||||
if terr.Hint != "" {
|
||||
h := terr.Hint
|
||||
hint = &h
|
||||
}
|
||||
} else {
|
||||
e := err.Error()
|
||||
errMsg = &e
|
||||
}
|
||||
}
|
||||
if serr := r.db.UpdateBatchTarget(j.id, idx, now, status, resultStr, errMsg, hint); serr != nil {
|
||||
r.log.Warn("batch: store target", "job", j.id, "idx", idx, "err", serr)
|
||||
}
|
||||
if resultStr != nil {
|
||||
return len(*resultStr)
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (r *Runner) callDepth(live int) {
|
||||
if r.cfg.Hooks.Depth != nil {
|
||||
r.cfg.Hooks.Depth(live)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Runner) callInflight(n int) {
|
||||
if r.cfg.Hooks.Inflight != nil {
|
||||
r.cfg.Hooks.Inflight(n)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Runner) callTarget(err error, d time.Duration, resultBytes int) {
|
||||
if r.cfg.Hooks.Target != nil {
|
||||
r.cfg.Hooks.Target(err, d, resultBytes)
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Runner) callJobSeconds(d time.Duration) {
|
||||
if r.cfg.Hooks.JobSeconds != nil {
|
||||
r.cfg.Hooks.JobSeconds(d)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
package config
|
||||
|
||||
import "testing"
|
||||
|
||||
// P6 owner review: a size too large for int64 falls back to the default instead of wrapping negative.
|
||||
func TestParseSizeOverflow(t *testing.T) {
|
||||
if got := ParseSize("9223372036854775807GiB", 7); got != 7 {
|
||||
t.Fatalf("ParseSize(huge) = %d, want the default", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
package config
|
||||
|
||||
// P6 task 08: the [batch] block and ParseSize.
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBatchBlock(t *testing.T) {
|
||||
p := filepath.Join(t.TempDir(), "sift.toml")
|
||||
body := "[batch]\ndisabled = false\ntargets_per_job = 16\njobs_per_caller = 2\njobs_total = 4\nworkers = 2\nttl = \"12h\"\nmax_bytes = \"32MiB\"\n"
|
||||
if err := os.WriteFile(p, []byte(body), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
f, err := Load(p)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
b := f.Batch
|
||||
if b.Disabled || b.TargetsPerJob != 16 || b.JobsPerCaller != 2 || b.JobsTotal != 4 || b.Workers != 2 || b.TTL != "12h" || b.MaxBytes != "32MiB" {
|
||||
t.Fatalf("batch block %+v", b)
|
||||
}
|
||||
// No block: zero values, and batch is on (the runner applies its own defaults).
|
||||
empty := filepath.Join(t.TempDir(), "empty.toml")
|
||||
if err := os.WriteFile(empty, []byte("listen = \":8443\"\n"), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if f, err := Load(empty); err != nil || f.Batch != (Batch{}) {
|
||||
t.Fatalf("no [batch] block: %+v %v", f.Batch, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseSize(t *testing.T) {
|
||||
const def = 7
|
||||
for _, c := range []struct {
|
||||
in string
|
||||
want int64
|
||||
}{
|
||||
{"64MiB", 64 << 20},
|
||||
{"512KiB", 512 << 10},
|
||||
{"1GiB", 1 << 30},
|
||||
{"1000", 1000},
|
||||
{" 2MiB ", 2 << 20},
|
||||
{"", def},
|
||||
{"junk", def},
|
||||
{"-1MiB", def},
|
||||
{"0", def},
|
||||
{"12MB", def}, // decimal units are not accepted: say MiB
|
||||
} {
|
||||
if got := ParseSize(c.in, def); got != c.want {
|
||||
t.Errorf("ParseSize(%q) = %d, want %d", c.in, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,10 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"math"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/pelletier/go-toml/v2"
|
||||
@@ -17,6 +20,19 @@ type File struct {
|
||||
Digest Digest `toml:"digest"`
|
||||
Context Context `toml:"context"`
|
||||
Search Search `toml:"search"`
|
||||
Batch Batch `toml:"batch"`
|
||||
}
|
||||
|
||||
// Batch is the [batch] block. Batch is on by default: no block means the runner runs with its own
|
||||
// defaults. disabled = true turns it off and the endpoints answer 503.
|
||||
type Batch struct {
|
||||
Disabled bool `toml:"disabled"`
|
||||
TargetsPerJob int `toml:"targets_per_job"`
|
||||
JobsPerCaller int `toml:"jobs_per_caller"`
|
||||
JobsTotal int `toml:"jobs_total"`
|
||||
Workers int `toml:"workers"`
|
||||
TTL string `toml:"ttl"`
|
||||
MaxBytes string `toml:"max_bytes"`
|
||||
}
|
||||
|
||||
type Dixie struct {
|
||||
@@ -110,3 +126,28 @@ func ParseDuration(s string, def time.Duration) time.Duration {
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
// ParseSize parses a positive byte count: a plain integer, or one with a binary suffix KiB, MiB or
|
||||
// GiB. Anything else, zero or negative returns def.
|
||||
func ParseSize(s string, def int64) int64 {
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "" {
|
||||
return def
|
||||
}
|
||||
var mult int64
|
||||
switch {
|
||||
case strings.HasSuffix(s, "KiB"):
|
||||
mult, s = 1<<10, s[:len(s)-3]
|
||||
case strings.HasSuffix(s, "MiB"):
|
||||
mult, s = 1<<20, s[:len(s)-3]
|
||||
case strings.HasSuffix(s, "GiB"):
|
||||
mult, s = 1<<30, s[:len(s)-3]
|
||||
default:
|
||||
mult = 1
|
||||
}
|
||||
n, err := strconv.ParseInt(s, 10, 64)
|
||||
if err != nil || n <= 0 || n > math.MaxInt64/mult { // the last: n*mult would overflow
|
||||
return def
|
||||
}
|
||||
return n * mult
|
||||
}
|
||||
|
||||
@@ -0,0 +1,342 @@
|
||||
package dcache
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
// terminal statuses mark a target as finished; a target already in one is not
|
||||
// counted again on a re-write.
|
||||
func isTerminal(status string) bool {
|
||||
return status == "done" || status == "failed"
|
||||
}
|
||||
|
||||
type BatchJob struct {
|
||||
ID string
|
||||
Caller string
|
||||
CallerClass string
|
||||
IDemKey string
|
||||
Status string
|
||||
Error string
|
||||
TargetsTotal int
|
||||
TargetsDone int
|
||||
CreatedAt int64
|
||||
Lease int64
|
||||
NBytes int64
|
||||
}
|
||||
|
||||
type BatchTarget struct {
|
||||
JobID string
|
||||
Idx int
|
||||
URL string
|
||||
Options string
|
||||
Status string
|
||||
Result *string
|
||||
Error *string
|
||||
Hint *string
|
||||
}
|
||||
|
||||
func nullString(s *string) sql.NullString {
|
||||
if s == nil {
|
||||
return sql.NullString{Valid: false}
|
||||
}
|
||||
return sql.NullString{String: *s, Valid: true}
|
||||
}
|
||||
|
||||
func nullPtr(n sql.NullString) *string {
|
||||
if !n.Valid {
|
||||
return nil
|
||||
}
|
||||
s := n.String
|
||||
return &s
|
||||
}
|
||||
|
||||
func (d *DB) CreateBatchJob(j BatchJob) error {
|
||||
if d == nil {
|
||||
return fmt.Errorf("dcache: nil db")
|
||||
}
|
||||
_, err := d.sql.Exec(`
|
||||
INSERT INTO batch_job
|
||||
(id, caller, caller_class, idem_key, status, error, targets_total, targets_done, created_at, lease, nbytes)
|
||||
VALUES(?,?,?,?,?,?,?,?,?,?,?)
|
||||
`, j.ID, j.Caller, j.CallerClass, j.IDemKey, j.Status, j.Error, j.TargetsTotal, j.TargetsDone, j.CreatedAt, j.Lease, j.NBytes)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *DB) BatchJob(id string) (BatchJob, bool) {
|
||||
if d == nil {
|
||||
return BatchJob{}, false
|
||||
}
|
||||
var j BatchJob
|
||||
err := d.sql.QueryRow(`
|
||||
SELECT id, caller, caller_class, idem_key, status, error,
|
||||
targets_total, targets_done, created_at, lease, nbytes
|
||||
FROM batch_job WHERE id = ?
|
||||
`, id).Scan(&j.ID, &j.Caller, &j.CallerClass, &j.IDemKey, &j.Status, &j.Error,
|
||||
&j.TargetsTotal, &j.TargetsDone, &j.CreatedAt, &j.Lease, &j.NBytes)
|
||||
if err != nil {
|
||||
d.mu.Lock()
|
||||
d.BatchM++
|
||||
d.mu.Unlock()
|
||||
return BatchJob{}, false
|
||||
}
|
||||
d.mu.Lock()
|
||||
d.BatchH++
|
||||
d.mu.Unlock()
|
||||
return j, true
|
||||
}
|
||||
|
||||
// BatchJobByIDem returns the id of the newest job for (caller, idemKey) that is
|
||||
// still within its lease and not failed/cancelled. Reads never write.
|
||||
func (d *DB) BatchJobByIDem(caller, idemKey string, now int64, ttl time.Duration) (string, bool) {
|
||||
if d == nil {
|
||||
return "", false
|
||||
}
|
||||
ttlSec := int64(ttl.Seconds())
|
||||
var id string
|
||||
err := d.sql.QueryRow(`
|
||||
SELECT id FROM batch_job
|
||||
WHERE caller = ? AND idem_key = ? AND lease + ? > ?
|
||||
AND status IN ('queued','running','done','partial')
|
||||
ORDER BY created_at DESC LIMIT 1
|
||||
`, caller, idemKey, ttlSec, now).Scan(&id)
|
||||
if err != nil {
|
||||
return "", false
|
||||
}
|
||||
return id, true
|
||||
}
|
||||
|
||||
func (d *DB) LiveBatchJobs(caller string) int {
|
||||
if d == nil {
|
||||
return 0
|
||||
}
|
||||
var n int
|
||||
_ = d.sql.QueryRow(`
|
||||
SELECT COUNT(*) FROM batch_job WHERE caller = ? AND status IN ('queued','running')
|
||||
`, caller).Scan(&n)
|
||||
return n
|
||||
}
|
||||
|
||||
func (d *DB) LiveBatchJobsTotal() int {
|
||||
if d == nil {
|
||||
return 0
|
||||
}
|
||||
var n int
|
||||
_ = d.sql.QueryRow(`SELECT COUNT(*) FROM batch_job WHERE status IN ('queued','running')`).Scan(&n)
|
||||
return n
|
||||
}
|
||||
|
||||
func (d *DB) PutBatchTargets(jobID string, targets []BatchTarget) error {
|
||||
if d == nil {
|
||||
return fmt.Errorf("dcache: nil db")
|
||||
}
|
||||
// One transaction: a job never has only some of its targets.
|
||||
tx, err := d.sql.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
for _, t := range targets {
|
||||
if _, err := tx.Exec(`
|
||||
INSERT INTO batch_target(job_id, idx, url, options, status, result, error, hint)
|
||||
VALUES(?,?,?,?,?,?,?,?)
|
||||
`, t.JobID, t.Idx, t.URL, t.Options, t.Status,
|
||||
nullString(t.Result), nullString(t.Error), nullString(t.Hint)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
func (d *DB) BatchTargets(jobID string) ([]BatchTarget, bool) {
|
||||
if d == nil {
|
||||
return nil, false
|
||||
}
|
||||
rows, err := d.sql.Query(`
|
||||
SELECT job_id, idx, url, options, status, result, error, hint
|
||||
FROM batch_target WHERE job_id = ? ORDER BY idx ASC
|
||||
`, jobID)
|
||||
if err != nil {
|
||||
d.mu.Lock()
|
||||
d.BatchM++
|
||||
d.mu.Unlock()
|
||||
return nil, false
|
||||
}
|
||||
defer rows.Close()
|
||||
var out []BatchTarget
|
||||
for rows.Next() {
|
||||
var (
|
||||
r, e, h sql.NullString
|
||||
t BatchTarget
|
||||
)
|
||||
if err := rows.Scan(&t.JobID, &t.Idx, &t.URL, &t.Options, &t.Status, &r, &e, &h); err != nil {
|
||||
d.mu.Lock()
|
||||
d.BatchM++
|
||||
d.mu.Unlock()
|
||||
return nil, false
|
||||
}
|
||||
t.Result = nullPtr(r)
|
||||
t.Error = nullPtr(e)
|
||||
t.Hint = nullPtr(h)
|
||||
out = append(out, t)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
d.mu.Lock()
|
||||
d.BatchM++
|
||||
d.mu.Unlock()
|
||||
return nil, false
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
return out, true
|
||||
}
|
||||
|
||||
func (d *DB) UpdateBatchTarget(jobID string, idx int, now int64, status string, result, errMsg, hint *string) error {
|
||||
if d == nil {
|
||||
return fmt.Errorf("dcache: nil db")
|
||||
}
|
||||
tx, err := d.sql.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
|
||||
var prev string
|
||||
err = tx.QueryRow(`SELECT status FROM batch_target WHERE job_id = ? AND idx = ?`, jobID, idx).Scan(&prev)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
var resBytes int64
|
||||
if result != nil {
|
||||
resBytes = int64(len(*result))
|
||||
}
|
||||
deltaDone := 0
|
||||
if isTerminal(status) && !isTerminal(prev) {
|
||||
deltaDone = 1
|
||||
}
|
||||
|
||||
if _, err := tx.Exec(`
|
||||
UPDATE batch_target SET status = ?, result = ?, error = ?, hint = ? WHERE job_id = ? AND idx = ?
|
||||
`, status, nullString(result), nullString(errMsg), nullString(hint), jobID, idx); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(`
|
||||
UPDATE batch_job SET lease = ?, targets_done = targets_done + ?, nbytes = nbytes + ? WHERE id = ?
|
||||
`, now, deltaDone, resBytes, jobID); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
func (d *DB) SetBatchJobStatus(id, status, errMsg string, now int64) error {
|
||||
if d == nil {
|
||||
return fmt.Errorf("dcache: nil db")
|
||||
}
|
||||
_, err := d.sql.Exec(`UPDATE batch_job SET status = ?, error = ?, lease = ? WHERE id = ?`, status, errMsg, now, id)
|
||||
return err
|
||||
}
|
||||
|
||||
func (d *DB) TouchBatchJob(id string, now int64) error {
|
||||
if d == nil {
|
||||
return fmt.Errorf("dcache: nil db")
|
||||
}
|
||||
_, err := d.sql.Exec(`UPDATE batch_job SET lease = ? WHERE id = ?`, now, id)
|
||||
return err
|
||||
}
|
||||
|
||||
func (d *DB) SweepBatch(now int64, ttl time.Duration) int {
|
||||
if d == nil {
|
||||
return 0
|
||||
}
|
||||
ttlSec := int64(ttl.Seconds())
|
||||
// A queued or running job is never swept: it is still in a worker's rotation. The restart
|
||||
// reaper marks such jobs failed, and they are swept after that. Both deletes are one
|
||||
// transaction, so a job never loses its targets but keeps its row.
|
||||
tx, err := d.sql.Begin()
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
const expired = `SELECT id FROM batch_job WHERE lease + ? <= ? AND status NOT IN ('queued','running')`
|
||||
if _, err := tx.Exec(`DELETE FROM batch_target WHERE job_id IN (`+expired+`)`, ttlSec, now); err != nil {
|
||||
return 0
|
||||
}
|
||||
res, err := tx.Exec(`DELETE FROM batch_job WHERE id IN (`+expired+`)`, ttlSec, now)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
if err := tx.Commit(); err != nil {
|
||||
return 0
|
||||
}
|
||||
n, _ := res.RowsAffected()
|
||||
return int(n)
|
||||
}
|
||||
|
||||
func (d *DB) EvictBatchOldest(maxBytes int64) int {
|
||||
if d == nil {
|
||||
return 0
|
||||
}
|
||||
evicted := 0
|
||||
for d.BatchNbytes() > maxBytes {
|
||||
var id string
|
||||
err := d.sql.QueryRow(`
|
||||
SELECT id FROM batch_job
|
||||
WHERE status NOT IN ('queued','running')
|
||||
ORDER BY lease ASC, id ASC
|
||||
LIMIT 1
|
||||
`).Scan(&id)
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
if err := d.deleteBatchJob(id); err != nil {
|
||||
break
|
||||
}
|
||||
evicted++
|
||||
}
|
||||
return evicted
|
||||
}
|
||||
|
||||
func (d *DB) MarkInFlightFailed(reason string, now int64) int {
|
||||
if d == nil {
|
||||
return 0
|
||||
}
|
||||
res, err := d.sql.Exec(`
|
||||
UPDATE batch_job SET status = 'failed', error = ?, lease = ? WHERE status IN ('queued','running')
|
||||
`, reason, now)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
n, _ := res.RowsAffected()
|
||||
return int(n)
|
||||
}
|
||||
|
||||
func (d *DB) BatchNbytes() int64 {
|
||||
if d == nil {
|
||||
return 0
|
||||
}
|
||||
var n sql.NullInt64
|
||||
_ = d.sql.QueryRow(`SELECT IFNULL(SUM(nbytes),0) FROM batch_job`).Scan(&n)
|
||||
return n.Int64
|
||||
}
|
||||
|
||||
// deleteBatchJob removes a job and its targets in one transaction.
|
||||
func (d *DB) deleteBatchJob(id string) error {
|
||||
tx, err := d.sql.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
if _, err := tx.Exec(`DELETE FROM batch_target WHERE job_id = ?`, id); err != nil {
|
||||
return err
|
||||
}
|
||||
if _, err := tx.Exec(`DELETE FROM batch_job WHERE id = ?`, id); err != nil {
|
||||
return err
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
package dcache
|
||||
|
||||
// P6 owner review: the sweep never removes a live job.
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestSweepSparesLiveJobs(t *testing.T) {
|
||||
db := newBatchDB(t)
|
||||
for _, st := range []string{"queued", "running", "done"} {
|
||||
if err := db.CreateBatchJob(BatchJob{ID: st, Caller: "a", CallerClass: ClassFleet, IDemKey: st, Status: st, TargetsTotal: 1, CreatedAt: 1, Lease: 1}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if n := db.SweepBatch(1_000_000, time.Hour); n != 1 {
|
||||
t.Fatalf("swept %d, want only the done job", n)
|
||||
}
|
||||
for _, id := range []string{"queued", "running"} {
|
||||
if _, ok := db.BatchJob(id); !ok {
|
||||
t.Fatalf("live job %s was swept", id)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,254 @@
|
||||
package dcache
|
||||
|
||||
import (
|
||||
"io"
|
||||
"log/slog"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func newBatchDB(t *testing.T) *DB {
|
||||
t.Helper()
|
||||
db, err := Open("", 1<<20, slog.New(slog.NewTextHandler(io.Discard, nil)))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
return db
|
||||
}
|
||||
|
||||
func strp(s string) *string { return &s }
|
||||
|
||||
func TestBatchJobRoundTrip(t *testing.T) {
|
||||
db := newBatchDB(t)
|
||||
j := BatchJob{
|
||||
ID: "j1", Caller: "alice", CallerClass: ClassFleet, IDemKey: "k1",
|
||||
Status: "queued", TargetsTotal: 3, CreatedAt: 100, Lease: 100,
|
||||
}
|
||||
if err := db.CreateBatchJob(j); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, ok := db.BatchJob("j1")
|
||||
if !ok {
|
||||
t.Fatal("missing job")
|
||||
}
|
||||
if got.ID != j.ID || got.Caller != "alice" || got.CallerClass != ClassFleet ||
|
||||
got.Status != "queued" || got.TargetsTotal != 3 || got.Lease != 100 || got.CreatedAt != 100 {
|
||||
t.Fatalf("round trip mismatch: %+v", got)
|
||||
}
|
||||
if _, ok := db.BatchJob("nope"); ok {
|
||||
t.Fatal("unknown id must be ok=false")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchIDemLookup(t *testing.T) {
|
||||
db := newBatchDB(t)
|
||||
if err := db.CreateBatchJob(BatchJob{ID: "j1", Caller: "a", CallerClass: ClassFleet, IDemKey: "k", Status: "running", TargetsTotal: 1, CreatedAt: 100, Lease: 1000}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if id, ok := db.BatchJobByIDem("a", "k", 2000, time.Hour); !ok || id != "j1" {
|
||||
t.Fatalf("live idem: %q %v", id, ok)
|
||||
}
|
||||
if _, ok := db.BatchJobByIDem("a", "k", 5000, time.Hour); ok {
|
||||
t.Fatal("lease-expired job must not match")
|
||||
}
|
||||
if _, ok := db.BatchJobByIDem("b", "k", 2000, time.Hour); ok {
|
||||
t.Fatal("other caller must not match")
|
||||
}
|
||||
// Only a job that is live or has results is reused. A failed (e.g. killed by a restart)
|
||||
// or cancelled job must not be handed back to a resubmit for the rest of its lease.
|
||||
for _, c := range []struct {
|
||||
status string
|
||||
reuse bool
|
||||
}{{"queued", true}, {"done", true}, {"partial", true}, {"failed", false}, {"cancelled", false}} {
|
||||
key := "k-" + c.status
|
||||
if err := db.CreateBatchJob(BatchJob{ID: "j-" + c.status, Caller: "a", CallerClass: ClassFleet, IDemKey: key, Status: c.status, TargetsTotal: 1, CreatedAt: 100, Lease: 1000}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, ok := db.BatchJobByIDem("a", key, 2000, time.Hour); ok != c.reuse {
|
||||
t.Errorf("status %s: reused=%v, want %v", c.status, ok, c.reuse)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchLiveCounts(t *testing.T) {
|
||||
db := newBatchDB(t)
|
||||
mk := func(id, caller, status string) {
|
||||
if err := db.CreateBatchJob(BatchJob{ID: id, Caller: caller, CallerClass: ClassFleet, IDemKey: id, Status: status, TargetsTotal: 1, CreatedAt: 1, Lease: 1}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
mk("j1", "a", "queued")
|
||||
mk("j2", "a", "running")
|
||||
mk("j3", "a", "done")
|
||||
mk("j4", "b", "running")
|
||||
if n := db.LiveBatchJobs("a"); n != 2 {
|
||||
t.Fatalf("caller a live=%d want 2", n)
|
||||
}
|
||||
if n := db.LiveBatchJobsTotal(); n != 3 {
|
||||
t.Fatalf("total live=%d want 3", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchTargetsOrderedAndUpdate(t *testing.T) {
|
||||
db := newBatchDB(t)
|
||||
must := func(err error) {
|
||||
t.Helper()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
must(db.CreateBatchJob(BatchJob{ID: "j1", Caller: "a", CallerClass: ClassFleet, IDemKey: "k", Status: "running", TargetsTotal: 3, CreatedAt: 100, Lease: 100}))
|
||||
must(db.PutBatchTargets("j1", []BatchTarget{
|
||||
{JobID: "j1", Idx: 0, URL: "https://a.example/", Options: `{"mode":"clean"}`, Status: "pending"},
|
||||
{JobID: "j1", Idx: 1, URL: "https://b.example/", Status: "pending"},
|
||||
{JobID: "j1", Idx: 2, URL: "https://c.example/", Status: "pending"},
|
||||
}))
|
||||
got, ok := db.BatchTargets("j1")
|
||||
if !ok || len(got) != 3 {
|
||||
t.Fatalf("targets len=%d ok=%v", len(got), ok)
|
||||
}
|
||||
for i := range got {
|
||||
if got[i].Idx != i {
|
||||
t.Fatalf("order broken: got[%d].Idx=%d", i, got[i].Idx)
|
||||
}
|
||||
}
|
||||
|
||||
// A done write bumps the lease to now and counts the target once.
|
||||
res := strp(`{"content":"ok"}`)
|
||||
must(db.UpdateBatchTarget("j1", 0, 200, "done", res, nil, nil))
|
||||
j, _ := db.BatchJob("j1")
|
||||
if j.Lease != 200 || j.TargetsDone != 1 {
|
||||
t.Fatalf("after first write lease=%d done=%d", j.Lease, j.TargetsDone)
|
||||
}
|
||||
// Re-writing the same target (still a done transition) must not double-count.
|
||||
must(db.UpdateBatchTarget("j1", 0, 201, "done", res, nil, nil))
|
||||
j, _ = db.BatchJob("j1")
|
||||
if j.TargetsDone != 1 {
|
||||
t.Fatalf("double-counted: done=%d", j.TargetsDone)
|
||||
}
|
||||
// A failed write counts too.
|
||||
must(db.UpdateBatchTarget("j1", 1, 202, "failed", nil, strp("upstream 500"), strp("browser")))
|
||||
j, _ = db.BatchJob("j1")
|
||||
if j.TargetsDone != 2 || j.Lease != 202 {
|
||||
t.Fatalf("after failed write lease=%d done=%d", j.Lease, j.TargetsDone)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchTouchAndStatus(t *testing.T) {
|
||||
db := newBatchDB(t)
|
||||
if err := db.CreateBatchJob(BatchJob{ID: "j1", Caller: "a", CallerClass: ClassFleet, IDemKey: "k", Status: "running", TargetsTotal: 1, CreatedAt: 100, Lease: 100}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.TouchBatchJob("j1", 500); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if j, _ := db.BatchJob("j1"); j.Lease != 500 {
|
||||
t.Fatalf("touch lease=%d want 500", j.Lease)
|
||||
}
|
||||
if err := db.SetBatchJobStatus("j1", "cancelled", "", 600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if j, _ := db.BatchJob("j1"); j.Status != "cancelled" || j.Lease != 600 {
|
||||
t.Fatalf("status=%q lease=%d", j.Status, j.Lease)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSweepBatch(t *testing.T) {
|
||||
db := newBatchDB(t)
|
||||
mk := func(id string, lease int64) {
|
||||
if err := db.CreateBatchJob(BatchJob{ID: id, Caller: "a", CallerClass: ClassFleet, IDemKey: id, Status: "done", TargetsTotal: 1, CreatedAt: 1, Lease: lease}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.PutBatchTargets(id, []BatchTarget{{JobID: id, Idx: 0, URL: "https://x.example/", Status: "done", Result: strp("x")}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
mk("old", 100)
|
||||
mk("new", 1000)
|
||||
n := db.SweepBatch(300, 100*time.Second) // old: 100+100=200<=300 expired; new: 1100>300 kept
|
||||
if n != 1 {
|
||||
t.Fatalf("swept %d want 1", n)
|
||||
}
|
||||
if _, ok := db.BatchJob("old"); ok {
|
||||
t.Fatal("old job must be gone")
|
||||
}
|
||||
if _, ok := db.BatchJob("new"); !ok {
|
||||
t.Fatal("new job must remain")
|
||||
}
|
||||
if _, ok := db.BatchTargets("old"); ok {
|
||||
t.Fatal("old targets must be gone")
|
||||
}
|
||||
}
|
||||
|
||||
func TestEvictBatchOldest(t *testing.T) {
|
||||
db := newBatchDB(t)
|
||||
done := func(id string, lease int64, size int) {
|
||||
if err := db.CreateBatchJob(BatchJob{ID: id, Caller: "a", CallerClass: ClassFleet, IDemKey: id, Status: "done", TargetsTotal: 1, CreatedAt: 1, Lease: lease}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.PutBatchTargets(id, []BatchTarget{{JobID: id, Idx: 0, URL: "https://x.example/", Status: "pending"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.UpdateBatchTarget(id, 0, lease, "done", strp(strings.Repeat("x", size)), nil, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
done("old1", 100, 100)
|
||||
done("old2", 200, 200)
|
||||
done("new", 300, 300)
|
||||
|
||||
// A running job with one completed target must never be byte-evicted.
|
||||
if err := db.CreateBatchJob(BatchJob{ID: "run", Caller: "a", CallerClass: ClassFleet, IDemKey: "run", Status: "running", TargetsTotal: 2, TargetsDone: 1, CreatedAt: 1, Lease: 400}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.PutBatchTargets("run", []BatchTarget{
|
||||
{JobID: "run", Idx: 0, URL: "https://x.example/", Status: "pending"},
|
||||
{JobID: "run", Idx: 1, URL: "https://y.example/", Status: "pending"},
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.UpdateBatchTarget("run", 0, 400, "done", strp(strings.Repeat("x", 500)), nil, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// total nbytes = 100+200+300+500 = 1100. Evict oldest done until <= 250:
|
||||
// old1(100), old2(200), new(300) all evicted (oldest first); run stays (live).
|
||||
db.EvictBatchOldest(250)
|
||||
for _, id := range []string{"old1", "old2", "new"} {
|
||||
if _, ok := db.BatchJob(id); ok {
|
||||
t.Fatalf("%s should be byte-evicted", id)
|
||||
}
|
||||
}
|
||||
if _, ok := db.BatchJob("run"); !ok {
|
||||
t.Fatal("running job must not be byte-evicted")
|
||||
}
|
||||
if n := db.BatchNbytes(); n != 500 {
|
||||
t.Fatalf("BatchNbytes=%d want 500", n)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMarkInFlightFailed(t *testing.T) {
|
||||
db := newBatchDB(t)
|
||||
mk := func(id, status string) {
|
||||
if err := db.CreateBatchJob(BatchJob{ID: id, Caller: "a", CallerClass: ClassFleet, IDemKey: id, Status: status, TargetsTotal: 1, CreatedAt: 1, Lease: 1}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
mk("q", "queued")
|
||||
mk("r", "running")
|
||||
mk("d", "done")
|
||||
if n := db.MarkInFlightFailed("restart", 999); n != 2 {
|
||||
t.Fatalf("marked %d want 2", n)
|
||||
}
|
||||
if j, _ := db.BatchJob("q"); j.Status != "failed" || j.Error != "restart" || j.Lease != 999 {
|
||||
t.Fatalf("q %+v", j)
|
||||
}
|
||||
if j, _ := db.BatchJob("r"); j.Status != "failed" || j.Error != "restart" {
|
||||
t.Fatalf("r %+v", j)
|
||||
}
|
||||
if j, _ := db.BatchJob("d"); j.Status != "done" {
|
||||
t.Fatalf("done job changed: %+v", j)
|
||||
}
|
||||
}
|
||||
+28
-13
@@ -11,8 +11,6 @@ import (
|
||||
"log/slog"
|
||||
"math"
|
||||
"os"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -47,6 +45,8 @@ type DB struct {
|
||||
SearchM uint64
|
||||
PDFH uint64
|
||||
PDFM uint64
|
||||
BatchH uint64
|
||||
BatchM uint64
|
||||
Evictions uint64
|
||||
}
|
||||
|
||||
@@ -110,6 +110,32 @@ func Open(path string, maxBytes int64, log *slog.Logger) (*DB, error) {
|
||||
text TEXT,
|
||||
PRIMARY KEY (k, page)
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS batch_job (
|
||||
id TEXT PRIMARY KEY,
|
||||
caller TEXT NOT NULL,
|
||||
caller_class TEXT NOT NULL,
|
||||
idem_key TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
error TEXT NOT NULL DEFAULT '',
|
||||
targets_total INTEGER NOT NULL,
|
||||
targets_done INTEGER NOT NULL DEFAULT 0,
|
||||
created_at INTEGER NOT NULL,
|
||||
lease INTEGER NOT NULL,
|
||||
nbytes INTEGER NOT NULL DEFAULT 0
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS batch_target (
|
||||
job_id TEXT NOT NULL,
|
||||
idx INTEGER NOT NULL,
|
||||
url TEXT NOT NULL,
|
||||
options TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
result TEXT,
|
||||
error TEXT,
|
||||
hint TEXT,
|
||||
PRIMARY KEY (job_id, idx)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS batch_job_idem ON batch_job(caller, idem_key);
|
||||
CREATE INDEX IF NOT EXISTS batch_job_lease ON batch_job(lease);
|
||||
`); err != nil {
|
||||
_ = sq.Close()
|
||||
return nil, err
|
||||
@@ -153,17 +179,6 @@ func QuestionEmbedKey(question, embedModel, callerClass string) string {
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
func SearchKey(query string, sources []string, k int, hour time.Time, callerClass string) string {
|
||||
if callerClass == "" {
|
||||
callerClass = ClassFleet
|
||||
}
|
||||
src := append([]string(nil), sources...)
|
||||
sort.Strings(src)
|
||||
bucket := hour.UTC().Format("2006-01-02T15")
|
||||
sum := sha256.Sum256([]byte(NormalizeQuestion(query) + "\x00" + strings.Join(src, ",") + "\x00" + strconv.Itoa(k) + "\x00" + bucket + "\x00" + callerClass))
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
type DigestHit struct {
|
||||
Text string
|
||||
ChunksUsed int
|
||||
|
||||
@@ -109,6 +109,12 @@ func TestSearchKeyPartitions(t *testing.T) {
|
||||
if SearchKey("hello", []string{"searxng", "tavily"}, 8, now.Add(time.Hour), ClassFleet) == a {
|
||||
t.Fatal("hour")
|
||||
}
|
||||
if EncodeSearchKey(SearchKeyInput{Queries: []string{"hello"}, Sources: []string{"searxng"}, K: 8, Hour: now, CallerClass: ClassFleet}) == EncodeSearchKey(SearchKeyInput{Queries: []string{"hello"}, Sources: []string{"searxng"}, K: 8, Hour: now, CallerClass: ClassFleet, IncludeDomains: []string{"a.com"}}) {
|
||||
t.Fatal("include_domains")
|
||||
}
|
||||
if EncodeSearchKey(SearchKeyInput{Queries: []string{"hello"}, Sources: []string{"searxng"}, K: 8, Hour: now, CallerClass: ClassFleet}) != EncodeSearchKey(SearchKeyInput{Queries: []string{"hello"}, Sources: []string{"searxng"}, K: 8, Hour: now, CallerClass: ClassFleet, Topic: "general"}) {
|
||||
t.Fatal("topic general")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCorruptDBBypass(t *testing.T) {
|
||||
|
||||
@@ -0,0 +1,108 @@
|
||||
package dcache
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"sort"
|
||||
"strconv"
|
||||
"time"
|
||||
)
|
||||
|
||||
// SearchKeyInput is the canonical cache identity for a search. Callers pass
|
||||
// already-normalised filter fields (long time_range, sorted domains optional).
|
||||
type SearchKeyInput struct {
|
||||
Queries []string
|
||||
Sources []string
|
||||
K int
|
||||
Hour time.Time
|
||||
CallerClass string
|
||||
IncludeDomains []string
|
||||
ExcludeDomains []string
|
||||
IncludeDomainsMode string
|
||||
TimeRange string
|
||||
StartDate string
|
||||
EndDate string
|
||||
Topic string
|
||||
ExactMatch bool
|
||||
MinScore float64
|
||||
}
|
||||
|
||||
func SearchKey(query string, sources []string, k int, hour time.Time, callerClass string) string {
|
||||
return EncodeSearchKey(SearchKeyInput{
|
||||
Queries: []string{query},
|
||||
Sources: sources,
|
||||
K: k,
|
||||
Hour: hour,
|
||||
CallerClass: callerClass,
|
||||
})
|
||||
}
|
||||
|
||||
func EncodeSearchKey(in SearchKeyInput) string {
|
||||
if in.CallerClass == "" {
|
||||
in.CallerClass = ClassFleet
|
||||
}
|
||||
src := append([]string(nil), in.Sources...)
|
||||
sort.Strings(src)
|
||||
qs := make([]string, len(in.Queries))
|
||||
for i, q := range in.Queries {
|
||||
qs[i] = NormalizeQuestion(q)
|
||||
}
|
||||
topic := in.Topic
|
||||
if topic == "" {
|
||||
topic = "general"
|
||||
}
|
||||
m := map[string]any{
|
||||
"caller_class": in.CallerClass,
|
||||
"exact_match": in.ExactMatch,
|
||||
"hour": in.Hour.UTC().Format("2006-01-02T15"),
|
||||
"k": in.K,
|
||||
"min_score": strconv.FormatFloat(in.MinScore, 'f', -1, 64),
|
||||
"queries": qs,
|
||||
"sources": src,
|
||||
"topic": topic,
|
||||
}
|
||||
if inc := sortedUnique(in.IncludeDomains); len(inc) > 0 {
|
||||
m["include_domains"] = inc
|
||||
mode := in.IncludeDomainsMode
|
||||
if mode == "" {
|
||||
mode = "restrict"
|
||||
}
|
||||
m["include_domains_mode"] = mode
|
||||
}
|
||||
if exc := sortedUnique(in.ExcludeDomains); len(exc) > 0 {
|
||||
m["exclude_domains"] = exc
|
||||
}
|
||||
if in.TimeRange != "" {
|
||||
m["time_range"] = in.TimeRange
|
||||
}
|
||||
if in.StartDate != "" {
|
||||
m["start_date"] = in.StartDate
|
||||
}
|
||||
if in.EndDate != "" {
|
||||
m["end_date"] = in.EndDate
|
||||
}
|
||||
b, _ := json.Marshal(m)
|
||||
sum := sha256.Sum256(b)
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
func sortedUnique(in []string) []string {
|
||||
if len(in) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := append([]string(nil), in...)
|
||||
sort.Strings(out)
|
||||
n := 0
|
||||
for _, s := range out {
|
||||
if s == "" {
|
||||
continue
|
||||
}
|
||||
if n > 0 && out[n-1] == s {
|
||||
continue
|
||||
}
|
||||
out[n] = s
|
||||
n++
|
||||
}
|
||||
return out[:n]
|
||||
}
|
||||
@@ -0,0 +1,172 @@
|
||||
package digest
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
)
|
||||
|
||||
// maxBgQueue is the number of documents the background worker may hold at once.
|
||||
const maxBgQueue = 4
|
||||
|
||||
// bgWorker is the single background goroutine that embeds the chunks of a long document after a
|
||||
// bm25 answer queued it. It owns a bounded, deduplicated queue and yields the Dixie gate to
|
||||
// interactive digests one batch at a time.
|
||||
type bgWorker struct {
|
||||
engine *Engine
|
||||
|
||||
mu sync.Mutex
|
||||
cond *sync.Cond
|
||||
queue []embedJob
|
||||
seen map[string]struct{}
|
||||
inProg bool
|
||||
started bool
|
||||
}
|
||||
|
||||
type embedJob struct {
|
||||
chunks []string
|
||||
contentHash string
|
||||
callerClass string
|
||||
}
|
||||
|
||||
// worker returns the engine's background worker, creating it once. It does not start the goroutine.
|
||||
func (e *Engine) worker() *bgWorker {
|
||||
e.bgOnce.Do(func() {
|
||||
w := &bgWorker{engine: e}
|
||||
w.cond = sync.NewCond(&w.mu)
|
||||
w.seen = map[string]struct{}{}
|
||||
e.bg = w
|
||||
})
|
||||
return e.bg
|
||||
}
|
||||
|
||||
// queueForEmbed asks the background worker to embed a document's chunks. The document is
|
||||
// deduplicated by contentHash+"|"+callerClass; a full queue drops it (and logs) so a later
|
||||
// question re-queues it.
|
||||
func (e *Engine) queueForEmbed(chunks []string, contentHash, callerClass string) {
|
||||
if e == nil {
|
||||
return
|
||||
}
|
||||
w := e.worker()
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
if !w.started {
|
||||
w.started = true
|
||||
go w.run()
|
||||
}
|
||||
key := contentHash + "|" + callerClass
|
||||
if _, ok := w.seen[key]; ok {
|
||||
return
|
||||
}
|
||||
if len(w.queue) >= maxBgQueue {
|
||||
// Leave it un-seen: a later question re-queues it.
|
||||
if e.Log != nil {
|
||||
e.Log.Warn("background embed queue full")
|
||||
}
|
||||
return
|
||||
}
|
||||
w.seen[key] = struct{}{}
|
||||
w.queue = append(w.queue, embedJob{chunks: chunks, contentHash: contentHash, callerClass: callerClass})
|
||||
w.cond.Signal()
|
||||
}
|
||||
|
||||
func (w *bgWorker) run() {
|
||||
for {
|
||||
w.mu.Lock()
|
||||
for len(w.queue) == 0 {
|
||||
w.cond.Wait()
|
||||
}
|
||||
job := w.queue[0]
|
||||
w.queue = w.queue[1:]
|
||||
w.inProg = true
|
||||
w.mu.Unlock()
|
||||
|
||||
w.process(job)
|
||||
|
||||
// Forget the document whether it finished or not: dedupe covers only queued and
|
||||
// in-progress work, so a later question re-queues a document whose embedding failed.
|
||||
w.mu.Lock()
|
||||
delete(w.seen, job.contentHash+"|"+job.callerClass)
|
||||
w.inProg = false
|
||||
w.mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
// process embeds the first MaxDocChunks chunks that lack a vector, EmbedBatch at a time, yielding
|
||||
// the gate between batches and storing each vector. On any embedding error it drops the rest of
|
||||
// the document.
|
||||
func (w *bgWorker) process(job embedJob) {
|
||||
e := w.engine
|
||||
key := func(i int) string {
|
||||
return dcache.EmbedKey(job.contentHash+"/sp1200", i, e.embedModel(), job.callerClass)
|
||||
}
|
||||
var idx []int
|
||||
for i := range job.chunks {
|
||||
if i >= MaxDocChunks {
|
||||
break
|
||||
}
|
||||
if _, ok := e.Cache.GetEmbed(key(i)); ok {
|
||||
continue
|
||||
}
|
||||
idx = append(idx, i)
|
||||
if len(idx) >= MaxDocChunks {
|
||||
break
|
||||
}
|
||||
}
|
||||
for b := 0; b < len(idx); b += EmbedBatch {
|
||||
j := b + EmbedBatch
|
||||
if j > len(idx) {
|
||||
j = len(idx)
|
||||
}
|
||||
batch := make([]string, j-b)
|
||||
for k := range batch {
|
||||
batch[k] = job.chunks[idx[b+k]]
|
||||
}
|
||||
w.acquireGate()
|
||||
ectx, cancel := context.WithTimeout(context.Background(), e.timeout())
|
||||
vecs, err := e.embed(ectx, batch)
|
||||
cancel()
|
||||
e.Gate.Release()
|
||||
if err != nil {
|
||||
if e.Log != nil {
|
||||
e.Log.Warn("background embed", "err", truncateErr(err, 512))
|
||||
}
|
||||
return
|
||||
}
|
||||
for k := range batch {
|
||||
if k < len(vecs) && len(vecs[k]) > 0 {
|
||||
e.Cache.PutEmbed(key(idx[b+k]), vecs[k])
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// acquireGate takes the Dixie gate, sleeping a second after each refusal (queue full or wait
|
||||
// expired) so interactive digests get their turns. Background work has no deadline: it waits.
|
||||
func (w *bgWorker) acquireGate() {
|
||||
for {
|
||||
if ok, _ := w.engine.Gate.Acquire(context.Background()); ok {
|
||||
return
|
||||
}
|
||||
time.Sleep(time.Second)
|
||||
}
|
||||
}
|
||||
|
||||
// WaitBackground returns when the background worker is idle: the queue empty and nothing embedding.
|
||||
func (e *Engine) WaitBackground(ctx context.Context) error {
|
||||
w := e.worker()
|
||||
for {
|
||||
w.mu.Lock()
|
||||
idle := len(w.queue) == 0 && !w.inProg
|
||||
w.mu.Unlock()
|
||||
if idle {
|
||||
return nil
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,109 @@
|
||||
package digest
|
||||
|
||||
import (
|
||||
"math"
|
||||
"strings"
|
||||
"unicode"
|
||||
)
|
||||
|
||||
const (
|
||||
bm25K1 = 1.2
|
||||
bm25B = 0.75
|
||||
bm25Min = 2
|
||||
)
|
||||
|
||||
// BM25 is a lexical ranker over a fixed set of documents. Terms are the same tokens as
|
||||
// terms() — lower-case runs of letters and digits, at least bm25Min characters — counted
|
||||
// with frequency rather than kept as a set.
|
||||
type BM25 struct {
|
||||
nDocs int
|
||||
avgDL float64
|
||||
df map[string]int
|
||||
toks []map[string]int
|
||||
len []int
|
||||
}
|
||||
|
||||
// NewBM25 indexes docs with the default BM25 parameters (k1 = 1.2, b = 0.75).
|
||||
func NewBM25(docs []string) *BM25 {
|
||||
b := &BM25{
|
||||
nDocs: len(docs),
|
||||
df: map[string]int{},
|
||||
toks: make([]map[string]int, len(docs)),
|
||||
len: make([]int, len(docs)),
|
||||
}
|
||||
var total int
|
||||
for i, doc := range docs {
|
||||
counts := docTerms(doc)
|
||||
n := 0
|
||||
for _, c := range counts {
|
||||
n += c
|
||||
}
|
||||
b.toks[i] = counts
|
||||
b.len[i] = n
|
||||
total += n
|
||||
for t := range counts {
|
||||
b.df[t]++
|
||||
}
|
||||
}
|
||||
if len(docs) > 0 {
|
||||
b.avgDL = float64(total) / float64(len(docs))
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// Scores returns one score per document, in order, each >= 0 and never NaN or Inf. An empty
|
||||
// query, an empty index or a document without any query term scores 0.
|
||||
func (b *BM25) Scores(query string) []float64 {
|
||||
if b == nil {
|
||||
return nil
|
||||
}
|
||||
if b.nDocs == 0 {
|
||||
return make([]float64, len(b.toks))
|
||||
}
|
||||
out := make([]float64, len(b.toks))
|
||||
qterms := docTerms(query)
|
||||
if len(qterms) == 0 {
|
||||
return out
|
||||
}
|
||||
avgDL := b.avgDL
|
||||
if avgDL == 0 {
|
||||
avgDL = 1
|
||||
}
|
||||
for i, toks := range b.toks {
|
||||
var score float64
|
||||
for term := range qterms {
|
||||
tf, ok := toks[term]
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
idf := math.Log(1 + (float64(b.nDocs)-float64(b.df[term])+0.5)/(float64(b.df[term])+0.5))
|
||||
denom := float64(tf) + bm25K1*(1-bm25B+bm25B*float64(b.len[i])/avgDL)
|
||||
score += idf * (float64(tf) * (bm25K1 + 1)) / denom
|
||||
}
|
||||
out[i] = score
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// docTerms counts occurrences of each BM25 term in s, using the same tokenization as terms().
|
||||
func docTerms(s string) map[string]int {
|
||||
out := map[string]int{}
|
||||
var b strings.Builder
|
||||
flush := func() {
|
||||
t := b.String()
|
||||
b.Reset()
|
||||
if len(t) < bm25Min {
|
||||
return
|
||||
}
|
||||
out[t]++
|
||||
}
|
||||
for _, r := range strings.ToLower(s) {
|
||||
if unicode.IsLetter(r) || unicode.IsDigit(r) {
|
||||
b.WriteRune(r)
|
||||
} else {
|
||||
flush()
|
||||
}
|
||||
}
|
||||
flush()
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
package digest
|
||||
|
||||
// P5b task 01: BM25, the lexical ranking used while a long document's embeddings are still
|
||||
// being built (and inside the hybrid score after).
|
||||
|
||||
import (
|
||||
"math"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBM25RanksTheRareTerm(t *testing.T) {
|
||||
docs := []string{
|
||||
"the ship sailed from helsinki to tallinn in the morning",
|
||||
"the ship stayed in port all day",
|
||||
"a zanzibarite sample was found near the harbour",
|
||||
"the harbour master logged the ship",
|
||||
}
|
||||
s := NewBM25(docs).Scores("where was the zanzibarite found")
|
||||
if len(s) != len(docs) {
|
||||
t.Fatalf("%d scores for %d docs", len(s), len(docs))
|
||||
}
|
||||
best := 0
|
||||
for i := range s {
|
||||
if s[i] > s[best] {
|
||||
best = i
|
||||
}
|
||||
}
|
||||
if best != 2 {
|
||||
t.Errorf("best = doc %d (%v), want doc 2, the only one with the rare term", best, s)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBM25Properties(t *testing.T) {
|
||||
filler := strings.Repeat("word ", 50)
|
||||
docs := []string{
|
||||
"alpha " + filler, // 0: tf 1
|
||||
strings.Repeat("alpha ", 10) + filler, // 1: tf 10, same filler
|
||||
"alpha common", // 2: tf 1, short
|
||||
"common " + filler, // 3: no alpha
|
||||
}
|
||||
idx := NewBM25(docs)
|
||||
s := idx.Scores("alpha")
|
||||
if !(s[1] > s[0]) {
|
||||
t.Errorf("more occurrences must score higher: tf10 %v vs tf1 %v", s[1], s[0])
|
||||
}
|
||||
if !(s[1] < 10*s[0]) {
|
||||
t.Errorf("term frequency must saturate: tf10 %v, 10 × tf1 %v", s[1], 10*s[0])
|
||||
}
|
||||
if !(s[2] > s[0]) {
|
||||
t.Errorf("the same term in a shorter document must score higher: short %v, long %v", s[2], s[0])
|
||||
}
|
||||
if s[3] != 0 {
|
||||
t.Errorf("a document without the term scores %v, want 0", s[3])
|
||||
}
|
||||
// IDF: a term in every document is worth less than a rare one.
|
||||
all := NewBM25([]string{"xx yy", "xx zz", "xx ww"}).Scores("xx")
|
||||
rare := NewBM25([]string{"xx yy", "qq zz", "qq ww"}).Scores("xx")
|
||||
if !(rare[0] > all[0]) {
|
||||
t.Errorf("idf: a term in 1 of 3 docs must outscore one in 3 of 3: %v vs %v", rare[0], all[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestBM25Degenerate(t *testing.T) {
|
||||
for _, docs := range [][]string{nil, {}, {""}, {"", "a b"}} {
|
||||
s := NewBM25(docs).Scores("a")
|
||||
if len(s) != len(docs) {
|
||||
t.Errorf("%q: %d scores", docs, len(s))
|
||||
}
|
||||
for _, v := range s {
|
||||
if math.IsNaN(v) || math.IsInf(v, 0) || v < 0 {
|
||||
t.Errorf("%q: score %v", docs, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, v := range NewBM25([]string{"a b", "c d"}).Scores("") {
|
||||
if v != 0 {
|
||||
t.Errorf("empty query scored %v", v)
|
||||
}
|
||||
}
|
||||
// Deterministic.
|
||||
a := NewBM25([]string{"alpha beta", "beta gamma", "gamma alpha"}).Scores("alpha gamma")
|
||||
b := NewBM25([]string{"alpha beta", "beta gamma", "gamma alpha"}).Scores("alpha gamma")
|
||||
for i := range a {
|
||||
if a[i] != b[i] {
|
||||
t.Fatal("scores differ between identical runs")
|
||||
}
|
||||
}
|
||||
}
|
||||
+20
-56
@@ -33,6 +33,8 @@ const (
|
||||
ChunkTokens = 1200
|
||||
EmbedBatch = 16
|
||||
EmbedMaxRunes = 3000
|
||||
EmbedSyncMax = 48
|
||||
MaxDocChunks = 5000
|
||||
DefaultPermit = 2
|
||||
DefaultQDepth = 8
|
||||
DefaultQWait = 20 * time.Second
|
||||
@@ -191,6 +193,9 @@ type Engine struct {
|
||||
Log *slog.Logger
|
||||
Cache *dcache.DB
|
||||
|
||||
bgOnce sync.Once
|
||||
bg *bgWorker
|
||||
|
||||
mu sync.Mutex
|
||||
Calls uint64
|
||||
Fallback map[string]uint64
|
||||
@@ -205,6 +210,8 @@ type Result struct {
|
||||
DigestMS int64
|
||||
ReturnedTok int
|
||||
Cached bool
|
||||
Retrieval string // "whole", "head", "hybrid", "bm25" or "cached"
|
||||
Embedded int // chunks that had an embedding when the answer was built
|
||||
}
|
||||
|
||||
func (e *Engine) timeout() time.Duration {
|
||||
@@ -305,6 +312,7 @@ func (e *Engine) Run(ctx context.Context, cleaned, question, finalURL, callerCla
|
||||
r.ChunksUsed = h.ChunksUsed
|
||||
r.ReturnedTok = EstTokens(h.Text)
|
||||
r.Cached = true
|
||||
r.Retrieval = "cached"
|
||||
r.DigestMS = time.Since(start).Milliseconds()
|
||||
return r
|
||||
}
|
||||
@@ -327,8 +335,10 @@ func (e *Engine) Run(ctx context.Context, cleaned, question, finalURL, callerCla
|
||||
dctx, cancel := context.WithTimeout(ctx, e.timeout())
|
||||
defer cancel()
|
||||
|
||||
body, nUsed, nTotal, err := e.selectText(dctx, cleaned, question)
|
||||
body, nUsed, nTotal, mode, embedded, err := e.selectText(dctx, cleaned, question, contentHash, callerClass)
|
||||
r.ChunksUsed, r.ChunksTotal = nUsed, nTotal
|
||||
r.Retrieval = mode
|
||||
r.Embedded = embedded
|
||||
if err != nil {
|
||||
e.logFallback("embed", err)
|
||||
r.Fallback = "embed"
|
||||
@@ -359,72 +369,26 @@ func (e *Engine) Run(ctx context.Context, cleaned, question, finalURL, callerCla
|
||||
}
|
||||
r.Text = text
|
||||
r.ReturnedTok = EstTokens(text)
|
||||
if e != nil && e.Cache != nil {
|
||||
// A bm25 answer is provisional: once the embeddings exist the same question
|
||||
// must get a hybrid answer, so it does not go in the digest cache.
|
||||
if e != nil && e.Cache != nil && mode != "bm25" {
|
||||
e.Cache.PutDigest(dkey, text, r.ChunksUsed)
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
func (e *Engine) selectText(ctx context.Context, cleaned, question string) (string, int, int, error) {
|
||||
func (e *Engine) selectText(ctx context.Context, cleaned, question, contentHash, callerClass string) (string, int, int, string, int, error) {
|
||||
if EstTokens(cleaned) <= SmallDoc {
|
||||
return cleaned, 1, 1, nil
|
||||
return cleaned, 1, 1, "whole", 0, nil
|
||||
}
|
||||
chunks := Split(cleaned, ChunkTokens)
|
||||
chunks := SplitPages(cleaned, ChunkTokens)
|
||||
if len(chunks) == 0 {
|
||||
return cleaned, 1, 1, nil
|
||||
return cleaned, 1, 1, "whole", 0, nil
|
||||
}
|
||||
if question == "" {
|
||||
return joinBudget(chunks, SmallDoc), min(len(chunks), countBudget(chunks, SmallDoc)), len(chunks), nil
|
||||
return joinBudget(chunks, SmallDoc), min(len(chunks), countBudget(chunks, SmallDoc)), len(chunks), "head", 0, nil
|
||||
}
|
||||
vecs, err := e.embed(ctx, append([]string{question}, capEmbed(chunks)...))
|
||||
if err != nil || len(vecs) < 2 {
|
||||
return joinBudget(chunks, SmallDoc), min(len(chunks), countBudget(chunks, SmallDoc)), len(chunks), err
|
||||
}
|
||||
q := vecs[0]
|
||||
type scored struct {
|
||||
i int
|
||||
s float64
|
||||
}
|
||||
var ss []scored
|
||||
for i, c := range chunks {
|
||||
if i+1 >= len(vecs) {
|
||||
break
|
||||
}
|
||||
ss = append(ss, scored{i: i, s: 0.7*cosine(q, vecs[i+1]) + 0.3*lexicalOverlap(question, c)})
|
||||
}
|
||||
picked := map[int]struct{}{}
|
||||
used := 0
|
||||
if len(chunks) > 0 {
|
||||
picked[0] = struct{}{}
|
||||
used = EstTokens(chunks[0])
|
||||
}
|
||||
for len(picked) < len(ss) {
|
||||
best, bi := -1.0, -1
|
||||
for _, x := range ss {
|
||||
if _, ok := picked[x.i]; ok {
|
||||
continue
|
||||
}
|
||||
if x.s > best {
|
||||
best, bi = x.s, x.i
|
||||
}
|
||||
}
|
||||
if bi < 0 {
|
||||
break
|
||||
}
|
||||
need := EstTokens(chunks[bi])
|
||||
if used+need > SmallDoc && used > 0 {
|
||||
break
|
||||
}
|
||||
picked[bi] = struct{}{}
|
||||
used += need
|
||||
}
|
||||
var ordered []string
|
||||
for i := range chunks {
|
||||
if _, ok := picked[i]; ok {
|
||||
ordered = append(ordered, chunks[i])
|
||||
}
|
||||
}
|
||||
return strings.Join(ordered, "\n\n"), len(ordered), len(chunks), nil
|
||||
return e.hybridOrBM25(ctx, chunks, question, contentHash, callerClass)
|
||||
}
|
||||
|
||||
func capEmbed(chunks []string) []string {
|
||||
|
||||
@@ -374,7 +374,7 @@ func TestPinChunk0Infobox(t *testing.T) {
|
||||
if EstTokens(cleaned) <= SmallDoc {
|
||||
t.Fatalf("fixture too small: %d", EstTokens(cleaned))
|
||||
}
|
||||
body, nUsed, nTotal, err := e.selectText(context.Background(), cleaned, "when was the initial release?")
|
||||
body, nUsed, nTotal, _, _, err := e.selectText(context.Background(), cleaned, "when was the initial release?", "", dcache.ClassFleet)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,149 @@
|
||||
package digest
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
)
|
||||
|
||||
// hybridOrBM25 scores a chunked document either hybrid (cosine + lexical overlap) or, when there
|
||||
// are too many uncached chunks to embed inline, lexical-only BM25. mode is "hybrid" or "bm25".
|
||||
func (e *Engine) hybridOrBM25(ctx context.Context, chunks []string, question, contentHash, callerClass string) (string, int, int, string, int, error) {
|
||||
key := func(i int) string {
|
||||
return dcache.EmbedKey(contentHash+"/sp1200", i, e.embedModel(), callerClass)
|
||||
}
|
||||
|
||||
if e.Cache == nil {
|
||||
if len(chunks) <= EmbedSyncMax {
|
||||
vecs, err := e.embed(ctx, append([]string{question}, capEmbed(chunks)...))
|
||||
if err != nil {
|
||||
return joinBudget(chunks, SmallDoc), min(len(chunks), countBudget(chunks, SmallDoc)), len(chunks), "hybrid", 0, err
|
||||
}
|
||||
chunkVecs := make([][]float64, len(chunks))
|
||||
for i := range chunks {
|
||||
chunkVecs[i] = vecs[i+1]
|
||||
}
|
||||
body, used := greedyBudget(chunks, hybridScores(chunks, question, vecs[0], chunkVecs), SmallDoc)
|
||||
return body, used, len(chunks), "hybrid", len(chunks), nil
|
||||
}
|
||||
body, used := greedyBudget(chunks, bm25Scores(chunks, question), SmallDoc)
|
||||
return body, used, len(chunks), "bm25", 0, nil
|
||||
}
|
||||
|
||||
// Cache present: look up every chunk's vector.
|
||||
var cached []int
|
||||
var cachedVecs [][]float64
|
||||
var uncached []int
|
||||
for i := range chunks {
|
||||
if v, ok := e.Cache.GetEmbed(key(i)); ok {
|
||||
cached = append(cached, i)
|
||||
cachedVecs = append(cachedVecs, v)
|
||||
continue
|
||||
}
|
||||
uncached = append(uncached, i)
|
||||
}
|
||||
|
||||
if len(uncached) <= EmbedSyncMax {
|
||||
in := append([]string{question}, capEmbed(uncachedChunks(chunks, uncached))...)
|
||||
vecs, err := e.embed(ctx, in)
|
||||
if err != nil {
|
||||
return joinBudget(chunks, SmallDoc), min(len(chunks), countBudget(chunks, SmallDoc)), len(chunks), "hybrid", len(cached), err
|
||||
}
|
||||
for j, i := range uncached {
|
||||
if len(vecs[j+1]) > 0 { // a vector the server left out is not cached as empty
|
||||
e.Cache.PutEmbed(key(i), vecs[j+1])
|
||||
}
|
||||
}
|
||||
chunkVecs := make([][]float64, len(chunks))
|
||||
for k, i := range cached {
|
||||
chunkVecs[i] = cachedVecs[k]
|
||||
}
|
||||
for j, i := range uncached {
|
||||
chunkVecs[i] = vecs[j+1]
|
||||
}
|
||||
body, used := greedyBudget(chunks, hybridScores(chunks, question, vecs[0], chunkVecs), SmallDoc)
|
||||
return body, used, len(chunks), "hybrid", len(chunks), nil
|
||||
}
|
||||
|
||||
body, used := greedyBudget(chunks, bm25Scores(chunks, question), SmallDoc)
|
||||
e.queueForEmbed(chunks, contentHash, callerClass)
|
||||
return body, used, len(chunks), "bm25", len(cached), nil
|
||||
}
|
||||
|
||||
func uncachedChunks(chunks []string, uncached []int) []string {
|
||||
out := make([]string, len(uncached))
|
||||
for j, i := range uncached {
|
||||
out[j] = chunks[i]
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// hybridScores builds the hybrid scorer 0.7*cosine(q, v_i) + 0.3*lexicalOverlap(question, chunk),
|
||||
// where chunkVecs[i] is chunk i's vector.
|
||||
func hybridScores(chunks []string, question string, q []float64, chunkVecs [][]float64) func(int) float64 {
|
||||
return func(i int) float64 {
|
||||
return 0.7*cosine(q, chunkVecs[i]) + 0.3*lexicalOverlap(question, chunks[i])
|
||||
}
|
||||
}
|
||||
|
||||
// bm25Scores builds the BM25 scorer, normalised by its maximum so scores land in [0, 1].
|
||||
func bm25Scores(chunks []string, question string) func(int) float64 {
|
||||
scores := NewBM25(chunks).Scores(question)
|
||||
var max float64
|
||||
for _, s := range scores {
|
||||
if s > max {
|
||||
max = s
|
||||
}
|
||||
}
|
||||
if max == 0 {
|
||||
return func(int) float64 { return 0 }
|
||||
}
|
||||
return func(i int) float64 { return scores[i] / max }
|
||||
}
|
||||
|
||||
// greedyBudget picks chunk 0 first, then the highest-scoring remaining chunks while the total
|
||||
// stays within budget tokens, and returns them in document order with the count used.
|
||||
func greedyBudget(chunks []string, score func(int) float64, budget int) (string, int) {
|
||||
type scored struct {
|
||||
i int
|
||||
s float64
|
||||
}
|
||||
var ss []scored
|
||||
for i := range chunks {
|
||||
ss = append(ss, scored{i: i, s: score(i)})
|
||||
}
|
||||
picked := map[int]struct{}{}
|
||||
used := 0
|
||||
if len(chunks) > 0 {
|
||||
picked[0] = struct{}{}
|
||||
used = EstTokens(chunks[0])
|
||||
}
|
||||
for len(picked) < len(ss) {
|
||||
best, bi := -1.0, -1
|
||||
for _, x := range ss {
|
||||
if _, ok := picked[x.i]; ok {
|
||||
continue
|
||||
}
|
||||
if x.s > best {
|
||||
best, bi = x.s, x.i
|
||||
}
|
||||
}
|
||||
if bi < 0 {
|
||||
break
|
||||
}
|
||||
need := EstTokens(chunks[bi])
|
||||
if used+need > budget && used > 0 {
|
||||
break
|
||||
}
|
||||
picked[bi] = struct{}{}
|
||||
used += need
|
||||
}
|
||||
var ordered []string
|
||||
for i := range chunks {
|
||||
if _, ok := picked[i]; ok {
|
||||
ordered = append(ordered, chunks[i])
|
||||
}
|
||||
}
|
||||
return strings.Join(ordered, "\n\n"), len(ordered)
|
||||
}
|
||||
@@ -0,0 +1,247 @@
|
||||
package digest
|
||||
|
||||
// P5b task 02: retrieval for long documents. Embeddings are cached per chunk; a document with
|
||||
// more uncached chunks than EmbedSyncMax is answered at once from BM25 while its chunks are
|
||||
// embedded in the background, and later questions use the hybrid score.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"hash/fnv"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.wntrmute.dev/kyle/sift/internal/dcache"
|
||||
)
|
||||
|
||||
// fakeDixie serves /embed (a bag-of-words hash into 64 dims, so cosine similarity means
|
||||
// something) and /chat (the SSE stub). It counts embedded inputs and can hold /embed shut.
|
||||
type fakeDixie struct {
|
||||
srv *httptest.Server
|
||||
embedded atomic.Int64 // inputs embedded, questions included
|
||||
hold chan struct{}
|
||||
holdOnce sync.Once
|
||||
lastChat atomic.Pointer[chatReq]
|
||||
}
|
||||
|
||||
func newFakeDixie(t *testing.T, hold bool) *fakeDixie {
|
||||
f := &fakeDixie{}
|
||||
if hold {
|
||||
f.hold = make(chan struct{})
|
||||
}
|
||||
var disc atomic.Bool
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/embed", func(w http.ResponseWriter, r *http.Request) {
|
||||
if f.hold != nil {
|
||||
select {
|
||||
case <-f.hold:
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
var body struct {
|
||||
Input []string `json:"input"`
|
||||
}
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
var data []map[string]any
|
||||
for i, s := range body.Input {
|
||||
data = append(data, map[string]any{"index": i, "embedding": bow(s)})
|
||||
}
|
||||
f.embedded.Add(int64(len(body.Input)))
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"data": data})
|
||||
})
|
||||
mux.Handle("/chat", sseHandler(t, "ok", &f.lastChat, &disc))
|
||||
f.srv = httptest.NewServer(mux)
|
||||
t.Cleanup(func() { f.release(); f.srv.Close() })
|
||||
return f
|
||||
}
|
||||
|
||||
func (f *fakeDixie) release() {
|
||||
if f.hold != nil {
|
||||
f.holdOnce.Do(func() { close(f.hold) })
|
||||
}
|
||||
}
|
||||
|
||||
func bow(s string) []float64 {
|
||||
v := make([]float64, 64)
|
||||
for t := range terms(s) {
|
||||
h := fnv.New32a()
|
||||
_, _ = h.Write([]byte(t))
|
||||
v[h.Sum32()%64]++
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
func (f *fakeDixie) engine(t *testing.T, gate *Gate, withCache bool) *Engine {
|
||||
e := &Engine{ChatURL: f.srv.URL + "/chat", EmbedURL: f.srv.URL + "/embed", Gate: gate, Client: f.srv.Client()}
|
||||
if withCache {
|
||||
db, err := dcache.Open(filepath.Join(t.TempDir(), "c.db"), 0, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
e.Cache = db
|
||||
}
|
||||
return e
|
||||
}
|
||||
|
||||
// longDoc is n pages of distinct filler; the page named rare carries a rare term.
|
||||
func longDoc(n, rare int) string {
|
||||
var b strings.Builder
|
||||
words := []string{"harbour", "vessel", "cargo", "anchor", "pilot", "tug", "berth", "quay", "hull", "keel", "mast", "deck"}
|
||||
for p := 1; p <= n; p++ {
|
||||
if p > 1 {
|
||||
b.WriteString("\n\n")
|
||||
}
|
||||
fmt.Fprintf(&b, "--- page %d ---\n\n", p)
|
||||
for l := 0; l < 40; l++ {
|
||||
fmt.Fprintf(&b, "page %d line %d %s %s %s %s\n", p, l, words[(p+l)%12], words[(p*3+l)%12], words[(p*5+l)%12], words[(p*7+l)%12])
|
||||
}
|
||||
if p == rare {
|
||||
b.WriteString("the zanzibarite cargo manifest was signed by the second officer\n")
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func chatDoc(f *fakeDixie) string {
|
||||
req := f.lastChat.Load()
|
||||
if req == nil {
|
||||
return ""
|
||||
}
|
||||
var s strings.Builder
|
||||
for _, m := range req.Messages {
|
||||
s.WriteString(m.Content)
|
||||
}
|
||||
return s.String()
|
||||
}
|
||||
|
||||
func TestEmbeddingsAreCachedAcrossQuestions(t *testing.T) {
|
||||
f := newFakeDixie(t, false)
|
||||
e := f.engine(t, NewGate(2, 8, 5*time.Second), true)
|
||||
doc := longDoc(30, 0) // ~11.5k tokens: over SmallDoc (so not "whole"), ~10 chunks: under EmbedSyncMax
|
||||
if EstTokens(doc) <= SmallDoc {
|
||||
t.Fatalf("test document is %d tokens; it must exceed SmallDoc (%d) to reach retrieval", EstTokens(doc), SmallDoc)
|
||||
}
|
||||
r := e.Run(context.Background(), doc, "which berth did the tug use?", "https://ex.test/a.pdf", dcache.ClassFleet)
|
||||
if r.Fallback != "" || r.Retrieval != "hybrid" {
|
||||
t.Fatalf("first question: fallback %q retrieval %q", r.Fallback, r.Retrieval)
|
||||
}
|
||||
n := len(SplitPages(doc, ChunkTokens))
|
||||
if got := f.embedded.Load(); got != int64(n+1) {
|
||||
t.Fatalf("first question embedded %d inputs, want %d chunks + the question", got, n)
|
||||
}
|
||||
r = e.Run(context.Background(), doc, "where was the pilot?", "https://ex.test/a.pdf", dcache.ClassFleet)
|
||||
if r.Fallback != "" || r.Retrieval != "hybrid" || r.Embedded != n {
|
||||
t.Fatalf("second question: %+v", r)
|
||||
}
|
||||
if got := f.embedded.Load(); got != int64(n+2) {
|
||||
t.Errorf("second question embedded %d inputs in total, want %d: only its own question is new", got, n+2)
|
||||
}
|
||||
}
|
||||
|
||||
// A long document is answered from BM25 without waiting for a single embedding, and becomes
|
||||
// hybrid once the background has embedded every chunk.
|
||||
func TestLongDocumentBM25ThenHybrid(t *testing.T) {
|
||||
f := newFakeDixie(t, true) // /embed held shut until release
|
||||
e := f.engine(t, NewGate(2, 8, 5*time.Second), true)
|
||||
doc := longDoc(300, 250)
|
||||
n := len(SplitPages(doc, ChunkTokens))
|
||||
if n <= EmbedSyncMax {
|
||||
t.Fatalf("test document has %d chunks; it must exceed EmbedSyncMax (%d)", n, EmbedSyncMax)
|
||||
}
|
||||
start := time.Now()
|
||||
r := e.Run(context.Background(), doc, "who signed the zanzibarite manifest?", "https://ex.test/big.pdf", dcache.ClassFleet)
|
||||
if d := time.Since(start); d > 10*time.Second {
|
||||
t.Fatalf("first answer took %v: it waited on embeddings", d)
|
||||
}
|
||||
if r.Fallback != "" || r.Retrieval != "bm25" || r.Embedded != 0 {
|
||||
t.Fatalf("first answer: %+v, want retrieval bm25 with 0 embedded", r)
|
||||
}
|
||||
if doc := chatDoc(f); !strings.Contains(doc, "zanzibarite") || !strings.Contains(doc, "page 250") {
|
||||
t.Errorf("BM25 did not bring page 250 (the only one with the term) into the prompt")
|
||||
}
|
||||
|
||||
f.release()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
||||
defer cancel()
|
||||
if err := e.WaitBackground(ctx); err != nil {
|
||||
t.Fatalf("background embedding did not finish: %v", err)
|
||||
}
|
||||
r = e.Run(context.Background(), doc, "who signed the zanzibarite manifest?", "https://ex.test/big.pdf", dcache.ClassFleet)
|
||||
if r.Fallback != "" || r.Retrieval != "hybrid" || r.Embedded != n {
|
||||
t.Fatalf("after background: %+v, want hybrid with all %d embedded", r, n)
|
||||
}
|
||||
if doc := chatDoc(f); !strings.Contains(doc, "zanzibarite") {
|
||||
t.Errorf("hybrid lost the page with the term")
|
||||
}
|
||||
}
|
||||
|
||||
// Two questions on the same long document queue it once: every chunk is embedded once.
|
||||
func TestBackgroundDedupes(t *testing.T) {
|
||||
f := newFakeDixie(t, true)
|
||||
e := f.engine(t, NewGate(2, 8, 5*time.Second), true)
|
||||
doc := longDoc(200, 0)
|
||||
n := len(SplitPages(doc, ChunkTokens))
|
||||
e.Run(context.Background(), doc, "first question about the quay", "https://ex.test/b.pdf", dcache.ClassFleet)
|
||||
e.Run(context.Background(), doc, "second question about the keel", "https://ex.test/b.pdf", dcache.ClassFleet)
|
||||
f.release()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
||||
defer cancel()
|
||||
if err := e.WaitBackground(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := f.embedded.Load(); got > int64(n) {
|
||||
t.Errorf("embedded %d inputs for a %d-chunk document: queued twice", got, n)
|
||||
}
|
||||
}
|
||||
|
||||
// Background work takes the gate one batch at a time, so an interactive digest is not starved.
|
||||
func TestBackgroundYieldsTheGate(t *testing.T) {
|
||||
f := newFakeDixie(t, false)
|
||||
e := f.engine(t, NewGate(1, 8, 10*time.Second), true)
|
||||
e.Run(context.Background(), longDoc(400, 0), "start the background", "https://ex.test/c.pdf", dcache.ClassFleet)
|
||||
start := time.Now()
|
||||
r := e.Run(context.Background(), "a short page about a tug and a quay", "what is here?", "https://ex.test/small", dcache.ClassFleet)
|
||||
if r.Fallback != "" {
|
||||
t.Fatalf("interactive digest fell back (%q) while background embedding ran", r.Fallback)
|
||||
}
|
||||
if d := time.Since(start); d > 5*time.Second {
|
||||
t.Errorf("interactive digest waited %v behind background embedding", d)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
|
||||
defer cancel()
|
||||
_ = e.WaitBackground(ctx)
|
||||
}
|
||||
|
||||
// Without a cache there is nowhere to keep background work: a long document is BM25 and
|
||||
// nothing is embedded.
|
||||
func TestNoCacheLongDocumentIsBM25Only(t *testing.T) {
|
||||
f := newFakeDixie(t, false)
|
||||
e := f.engine(t, NewGate(2, 8, 5*time.Second), false)
|
||||
r := e.Run(context.Background(), longDoc(300, 250), "who signed the zanzibarite manifest?", "https://ex.test/d.pdf", dcache.ClassFleet)
|
||||
if r.Fallback != "" || r.Retrieval != "bm25" {
|
||||
t.Fatalf("%+v, want bm25", r)
|
||||
}
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
if got := f.embedded.Load(); got != 0 {
|
||||
t.Errorf("embedded %d inputs with no cache to keep them in", got)
|
||||
}
|
||||
}
|
||||
|
||||
// A document that fits the digest budget whole is sent whole, as today.
|
||||
func TestSmallDocumentIsWhole(t *testing.T) {
|
||||
f := newFakeDixie(t, false)
|
||||
e := f.engine(t, NewGate(2, 8, 5*time.Second), true)
|
||||
r := e.Run(context.Background(), "--- page 1 ---\n\nOne short page.", "what is it?", "https://ex.test/e.pdf", dcache.ClassFleet)
|
||||
if r.Fallback != "" || r.Retrieval != "whole" {
|
||||
t.Fatalf("%+v, want whole", r)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package digest
|
||||
|
||||
// P5b owner review: regression tests for the background embedder.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// flakyDixie is fakeDixie whose /embed answers 500 to its first failN requests.
|
||||
func flakyDixie(t *testing.T, failN int64) (*fakeDixie, *atomic.Int64) {
|
||||
f := &fakeDixie{}
|
||||
var calls atomic.Int64
|
||||
var disc atomic.Bool
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/embed", func(w http.ResponseWriter, r *http.Request) {
|
||||
if calls.Add(1) <= failN {
|
||||
http.Error(w, "boom", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
Input []string `json:"input"`
|
||||
}
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
var data []map[string]any
|
||||
for i, s := range body.Input {
|
||||
data = append(data, map[string]any{"index": i, "embedding": bow(s)})
|
||||
}
|
||||
f.embedded.Add(int64(len(body.Input)))
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{"data": data})
|
||||
})
|
||||
mux.Handle("/chat", sseHandler(t, "ok", &f.lastChat, &disc))
|
||||
f.srv = httptest.NewServer(mux)
|
||||
t.Cleanup(f.srv.Close)
|
||||
return f, &calls
|
||||
}
|
||||
|
||||
// A document whose background embedding failed is queued again by a later question; it must
|
||||
// not stay "seen" (and so BM25-only) for the life of the process.
|
||||
func TestBackgroundRequeuesAfterFailure(t *testing.T) {
|
||||
f, _ := flakyDixie(t, 1)
|
||||
e := f.engine(t, NewGate(2, 8, 5*time.Second), true)
|
||||
doc := longDoc(300, 250)
|
||||
n := len(SplitPages(doc, ChunkTokens))
|
||||
wait := func() {
|
||||
t.Helper()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancel()
|
||||
if err := e.WaitBackground(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
q := "who signed the zanzibarite manifest?"
|
||||
if r := e.Run(context.Background(), doc, q, "https://ex.test/f.pdf", "fleet"); r.Retrieval != "bm25" {
|
||||
t.Fatalf("first: %+v", r)
|
||||
}
|
||||
wait() // the first batch failed: the document was dropped
|
||||
if r := e.Run(context.Background(), doc, q, "https://ex.test/f.pdf", "fleet"); r.Retrieval != "bm25" {
|
||||
t.Fatalf("second: %+v, want bm25 (nothing embedded yet)", r)
|
||||
}
|
||||
wait() // re-queued by the second question, and embedded this time
|
||||
r := e.Run(context.Background(), doc, q, "https://ex.test/f.pdf", "fleet")
|
||||
if r.Retrieval != "hybrid" || r.Embedded != n {
|
||||
t.Fatalf("third: %+v, want hybrid with all %d embedded", r, n)
|
||||
}
|
||||
}
|
||||
|
||||
// The background worker waits for a busy gate; it does not give up on the document.
|
||||
func TestBackgroundWaitsForBusyGate(t *testing.T) {
|
||||
f := newFakeDixie(t, false)
|
||||
gate := NewGate(1, 0, 50*time.Millisecond) // no queue: a busy gate refuses at once
|
||||
e := f.engine(t, gate, true)
|
||||
doc := longDoc(100, 0)
|
||||
n := len(SplitPages(doc, ChunkTokens))
|
||||
if ok, _ := gate.Acquire(context.Background()); !ok {
|
||||
t.Fatal("could not take the gate")
|
||||
}
|
||||
// Run takes the gate before it chooses a retrieval mode, so queue the document directly.
|
||||
chunks := SplitPages(doc, ChunkTokens)
|
||||
e.queueForEmbed(chunks, "hash-g", "fleet")
|
||||
time.Sleep(2500 * time.Millisecond) // longer than the old single retry
|
||||
gate.Release()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancel()
|
||||
if err := e.WaitBackground(ctx); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got := f.embedded.Load(); got != int64(n) {
|
||||
t.Fatalf("embedded %d of %d chunks: the worker gave up on a busy gate", got, n)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
package digest
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// pageMarker matches a page header line like "--- page 7 ---"; the number is group 1.
|
||||
var pageMarker = regexp.MustCompile(`^--- page (\d+) ---$`)
|
||||
|
||||
// SplitPages splits text into chunks of at most maxTok tokens (EstTokens) on line
|
||||
// boundaries, like Split, but page-aware. The text is expected to carry "--- page N ---"
|
||||
// markers (as pdfdoc.Full produces). A marker line sets the current page; a chunk that
|
||||
// starts after the first marker but not on a marker line is prefixed with
|
||||
// "--- page N (cont.) ---\n\n" for the page its first line belongs to, so every chunk
|
||||
// names the page it cites. Text with no marker lines yields exactly Split's result.
|
||||
func SplitPages(text string, maxTok int) []string {
|
||||
if maxTok <= 0 {
|
||||
maxTok = ChunkTokens
|
||||
}
|
||||
lines := strings.Split(text, "\n")
|
||||
var chunks []string
|
||||
var buf strings.Builder
|
||||
var firstLine string
|
||||
var chunkPage int
|
||||
hasLines, startedAfterMarker, markerSeen := false, false, false
|
||||
currentPage := 0
|
||||
|
||||
flush := func() {
|
||||
s := strings.TrimSpace(buf.String())
|
||||
if s == "" {
|
||||
buf.Reset()
|
||||
hasLines = false
|
||||
return
|
||||
}
|
||||
if pageMarker.MatchString(firstLine) || !startedAfterMarker {
|
||||
chunks = append(chunks, s)
|
||||
} else {
|
||||
chunks = append(chunks, fmt.Sprintf("--- page %d (cont.) ---\n\n", chunkPage)+s)
|
||||
}
|
||||
buf.Reset()
|
||||
hasLines = false
|
||||
}
|
||||
|
||||
for _, line := range lines {
|
||||
if m := pageMarker.FindStringSubmatch(line); m != nil {
|
||||
if n, err := strconv.Atoi(m[1]); err == nil {
|
||||
currentPage = n
|
||||
markerSeen = true
|
||||
}
|
||||
}
|
||||
next := line + "\n"
|
||||
if hasLines {
|
||||
est := EstTokens(buf.String() + next)
|
||||
if startedAfterMarker {
|
||||
est += EstTokens(fmt.Sprintf("--- page %d (cont.) ---\n\n", chunkPage))
|
||||
}
|
||||
if est > maxTok {
|
||||
flush()
|
||||
}
|
||||
}
|
||||
if !hasLines {
|
||||
firstLine = line
|
||||
chunkPage = currentPage
|
||||
startedAfterMarker = markerSeen
|
||||
hasLines = true
|
||||
}
|
||||
buf.WriteString(next)
|
||||
}
|
||||
flush()
|
||||
return chunks
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
package digest
|
||||
|
||||
// P5b task 01: SplitPages keeps every chunk of a PDF's text attached to its page, so a digest
|
||||
// built from a chunk that starts mid-page can still cite the page.
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// pagesText builds n pages of about perPage tokens each, every line naming its page.
|
||||
func pagesText(n, perPage int) string {
|
||||
var b strings.Builder
|
||||
for p := 1; p <= n; p++ {
|
||||
if p > 1 {
|
||||
b.WriteString("\n\n")
|
||||
}
|
||||
fmt.Fprintf(&b, "--- page %d ---\n\n", p)
|
||||
for l := 0; l*12 < perPage; l++ {
|
||||
fmt.Fprintf(&b, "page %d line %d alpha bravo charlie delta echo foxtrot golf hotel\n", p, l)
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
var markerRE = regexp.MustCompile(`^--- page (\d+)( \(cont\.\))? ---`)
|
||||
|
||||
func TestSplitPagesEveryChunkNamesItsPage(t *testing.T) {
|
||||
text := pagesText(4, 3000) // each page ~3000 tokens: several chunks per page at 1200
|
||||
chunks := SplitPages(text, ChunkTokens)
|
||||
if len(chunks) < 8 {
|
||||
t.Fatalf("%d chunks, want at least 8 for 4 × 3000 tokens at %d", len(chunks), ChunkTokens)
|
||||
}
|
||||
cont := 0
|
||||
for i, c := range chunks {
|
||||
m := markerRE.FindStringSubmatch(c)
|
||||
if m == nil {
|
||||
t.Fatalf("chunk %d does not start with a page marker: %.80q", i, c)
|
||||
}
|
||||
if m[2] != "" {
|
||||
cont++
|
||||
// A continuation marker names the page the chunk's first line belongs to.
|
||||
first := strings.SplitN(strings.TrimPrefix(c, m[0]), "\n", 3)
|
||||
line := strings.TrimSpace(strings.Join(first, " "))
|
||||
if !strings.Contains(line, "page "+m[1]+" line ") {
|
||||
t.Errorf("chunk %d marked page %s but starts with %.60q", i, m[1], line)
|
||||
}
|
||||
}
|
||||
if EstTokens(c) > ChunkTokens+20 {
|
||||
t.Errorf("chunk %d is %d tokens, over the %d budget (plus a marker line)", i, EstTokens(c), ChunkTokens)
|
||||
}
|
||||
}
|
||||
if cont == 0 {
|
||||
t.Error("no continuation markers: pages longer than a chunk must produce some")
|
||||
}
|
||||
}
|
||||
|
||||
// Nothing is lost or duplicated: removing the continuation markers gives back the original lines.
|
||||
func TestSplitPagesPreservesContent(t *testing.T) {
|
||||
text := pagesText(3, 2500)
|
||||
var got []string
|
||||
for _, c := range SplitPages(text, ChunkTokens) {
|
||||
for _, line := range strings.Split(c, "\n") {
|
||||
if strings.HasSuffix(line, " (cont.) ---") || strings.TrimSpace(line) == "" {
|
||||
continue
|
||||
}
|
||||
got = append(got, line)
|
||||
}
|
||||
}
|
||||
var want []string
|
||||
for _, line := range strings.Split(text, "\n") {
|
||||
if strings.TrimSpace(line) != "" {
|
||||
want = append(want, line)
|
||||
}
|
||||
}
|
||||
if strings.Join(got, "\n") != strings.Join(want, "\n") {
|
||||
t.Errorf("content changed: %d lines back, %d in", len(got), len(want))
|
||||
}
|
||||
}
|
||||
|
||||
// Text with no markers splits exactly as Split does.
|
||||
func TestSplitPagesWithoutMarkersIsSplit(t *testing.T) {
|
||||
text := strings.Repeat("plain words without any page markers at all\n", 2000)
|
||||
a, b := SplitPages(text, ChunkTokens), Split(text, ChunkTokens)
|
||||
if strings.Join(a, "\x00") != strings.Join(b, "\x00") {
|
||||
t.Errorf("SplitPages differs from Split on unmarked text: %d vs %d chunks", len(a), len(b))
|
||||
}
|
||||
}
|
||||
@@ -333,6 +333,21 @@ func ipv4Compatible(addr netip.Addr) (netip.Addr, bool) {
|
||||
return v4, true
|
||||
}
|
||||
|
||||
// PreCheck is CheckURL without name resolution: scheme, userinfo, host and port, and an IP
|
||||
// literal checked against the address rules. A hostname is not resolved here; the fetch-time
|
||||
// CheckURL and the dial-time Control still apply to it. The batch endpoint uses it to refuse a
|
||||
// target at submit time.
|
||||
func (g *Guard) PreCheck(raw string) error {
|
||||
_, literal, err := ParseURL(raw)
|
||||
if err != nil && !(literal.IsValid() && g.allows(literal)) {
|
||||
return err
|
||||
}
|
||||
if literal.IsValid() {
|
||||
return g.CheckAddr(literal)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// CheckURL parses, and if the host is a name, resolves every record. Any
|
||||
// blocked address fails the whole URL (mixed-record trick).
|
||||
func (g *Guard) CheckURL(ctx context.Context, raw string) (*url.URL, []netip.Addr, error) {
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
package guard
|
||||
|
||||
// P6 (owner, before the run): PreCheck is CheckURL without name resolution, for refusing a batch target at
|
||||
// submit time.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/netip"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestPreCheck(t *testing.T) {
|
||||
g := &Guard{
|
||||
AllowPrefixes: []netip.Prefix{netip.MustParsePrefix("127.0.0.0/8")},
|
||||
LookupIP: func(context.Context, string) ([]netip.Addr, error) {
|
||||
t.Fatal("PreCheck resolved a name")
|
||||
return nil, nil
|
||||
},
|
||||
}
|
||||
for _, u := range []string{
|
||||
"https://example.com/",
|
||||
"http://example.com/path?q=1",
|
||||
"https://internal.test/", // a name: not resolved here, so not refused here
|
||||
"http://127.0.0.1:8080/", // an allowed literal is exempt from the port rule, as in CheckURL
|
||||
"http://93.184.215.14/", // a public literal
|
||||
"http://0x7f000001:8080/", // inet_aton form of 127.0.0.1: an allowed literal
|
||||
} {
|
||||
if err := g.PreCheck(u); err != nil {
|
||||
t.Errorf("PreCheck(%q) = %v, want nil", u, err)
|
||||
}
|
||||
}
|
||||
for _, u := range []string{
|
||||
"",
|
||||
"ftp://example.com/",
|
||||
"http://user@example.com/",
|
||||
"http://example.com:8080/",
|
||||
"http://10.0.0.5/",
|
||||
"http://169.254.169.254/latest/meta-data/",
|
||||
"http://[::1]/",
|
||||
"http://0x0a000005/", // inet_aton form of 10.0.0.5
|
||||
} {
|
||||
err := g.PreCheck(u)
|
||||
if err == nil || !strings.HasPrefix(err.Error(), "blocked") {
|
||||
t.Errorf("PreCheck(%q) = %v, want a blocked error", u, err)
|
||||
}
|
||||
}
|
||||
// With no allow list, loopback is refused like any other private literal.
|
||||
var zero Guard
|
||||
if err := zero.PreCheck("http://127.0.0.1/"); err == nil {
|
||||
t.Error("a zero Guard allowed loopback")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
package metrics
|
||||
|
||||
import "time"
|
||||
|
||||
var batchResultBytesBuckets = []float64{1024, 4096, 16384, 65536, 262144, 1048576}
|
||||
|
||||
var batchJobOutcomes = map[string]struct{}{
|
||||
"created": {}, "reused": {}, "done": {}, "partial": {},
|
||||
"failed": {}, "cancelled": {}, "restart": {},
|
||||
}
|
||||
|
||||
var batchTargetOutcomes = map[string]struct{}{
|
||||
"ok": {}, "blocked": {}, "timeout": {}, "busy": {},
|
||||
"upstream_4xx": {}, "upstream_5xx": {}, "unsupported": {},
|
||||
"too_long": {}, "cancel": {},
|
||||
}
|
||||
|
||||
// ObserveBatchJob records one batch outcome. Values outside the known set fold
|
||||
// to "other" so a raw string never becomes a label.
|
||||
func (s *Set) ObserveBatchJob(outcome string) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
if _, ok := batchJobOutcomes[outcome]; !ok {
|
||||
outcome = "other"
|
||||
}
|
||||
s.BatchJobsTotal.WithLabelValues(outcome).Inc()
|
||||
}
|
||||
|
||||
// ObserveBatchTarget records one target outcome. Values outside the known set
|
||||
// fold to "other" so a raw string never becomes a label.
|
||||
func (s *Set) ObserveBatchTarget(outcome string) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
if _, ok := batchTargetOutcomes[outcome]; !ok {
|
||||
outcome = "other"
|
||||
}
|
||||
s.BatchTargetsTotal.WithLabelValues(outcome).Inc()
|
||||
}
|
||||
|
||||
// SetBatchQueueDepth sets the live job count gauge (queued+running).
|
||||
func (s *Set) SetBatchQueueDepth(n int) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.BatchQueueDepth.Set(float64(n))
|
||||
}
|
||||
|
||||
// SetBatchTargetsInflight sets the in-flight targets gauge.
|
||||
func (s *Set) SetBatchTargetsInflight(n int) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.BatchTargetsInflight.Set(float64(n))
|
||||
}
|
||||
|
||||
// ObserveBatchJobSeconds records the wall time from submit to a batch's terminal status.
|
||||
func (s *Set) ObserveBatchJobSeconds(d time.Duration) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.BatchJobSeconds.Observe(d.Seconds())
|
||||
}
|
||||
|
||||
// ObserveBatchTargetSeconds records the wall time for one target fetch.
|
||||
func (s *Set) ObserveBatchTargetSeconds(d time.Duration) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.BatchTargetSeconds.Observe(d.Seconds())
|
||||
}
|
||||
|
||||
// ObserveBatchResultBytes records the stored byte count of one target result.
|
||||
func (s *Set) ObserveBatchResultBytes(n int) {
|
||||
if s == nil {
|
||||
return
|
||||
}
|
||||
s.BatchResultBytes.Observe(float64(n))
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
package metrics
|
||||
|
||||
import (
|
||||
"io"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type batchSrc struct{}
|
||||
|
||||
func (batchSrc) Snapshot() Snapshot { return Snapshot{} }
|
||||
|
||||
func TestBatchSeries(t *testing.T) {
|
||||
s := New(batchSrc{})
|
||||
s.ObserveBatchJob("created")
|
||||
s.ObserveBatchJob("done")
|
||||
s.ObserveBatchJob("done")
|
||||
s.ObserveBatchJob("reused")
|
||||
s.ObserveBatchJob("some raw string") // folded to "other", never a label
|
||||
s.ObserveBatchTarget("ok")
|
||||
s.ObserveBatchTarget("ok")
|
||||
s.ObserveBatchTarget("blocked")
|
||||
s.ObserveBatchTarget("upstream_5xx")
|
||||
s.ObserveBatchTarget("leak me") // folded to "other"
|
||||
s.SetBatchQueueDepth(3)
|
||||
s.SetBatchTargetsInflight(1)
|
||||
s.ObserveBatchJobSeconds(4 * time.Second)
|
||||
s.ObserveBatchTargetSeconds(1500 * time.Millisecond)
|
||||
s.ObserveBatchResultBytes(4096)
|
||||
|
||||
req := httptest.NewRequest("GET", "/metrics", nil)
|
||||
req.RemoteAddr, req.Host = "127.0.0.1:9", "127.0.0.1:8890"
|
||||
rec := httptest.NewRecorder()
|
||||
s.Handler().ServeHTTP(rec, req)
|
||||
body, _ := io.ReadAll(rec.Body)
|
||||
out := string(body)
|
||||
for _, want := range []string{
|
||||
`sift_batch_jobs_total{outcome="created"} 1`,
|
||||
`sift_batch_jobs_total{outcome="done"} 2`,
|
||||
`sift_batch_jobs_total{outcome="reused"} 1`,
|
||||
`sift_batch_jobs_total{outcome="other"} 1`,
|
||||
`sift_batch_targets_total{outcome="ok"} 2`,
|
||||
`sift_batch_targets_total{outcome="blocked"} 1`,
|
||||
`sift_batch_targets_total{outcome="upstream_5xx"} 1`,
|
||||
`sift_batch_targets_total{outcome="other"} 1`,
|
||||
`sift_batch_queue_depth 3`,
|
||||
`sift_batch_targets_inflight 1`,
|
||||
`sift_batch_job_seconds_count 1`,
|
||||
`sift_batch_target_seconds_count 1`,
|
||||
`sift_batch_result_bytes_count 1`,
|
||||
} {
|
||||
if !strings.Contains(out, want) {
|
||||
t.Errorf("metrics lack %q", want)
|
||||
}
|
||||
}
|
||||
if strings.Contains(out, "raw") || strings.Contains(out, "leak") {
|
||||
t.Error("a raw outcome string leaked into a label")
|
||||
}
|
||||
}
|
||||
@@ -58,6 +58,13 @@ type Set struct {
|
||||
PDFPages prometheus.Counter
|
||||
PDFWait prometheus.Histogram
|
||||
PDFCacheHits prometheus.Counter
|
||||
BatchJobsTotal *prometheus.CounterVec
|
||||
BatchTargetsTotal *prometheus.CounterVec
|
||||
BatchQueueDepth prometheus.Gauge
|
||||
BatchTargetsInflight prometheus.Gauge
|
||||
BatchJobSeconds prometheus.Histogram
|
||||
BatchTargetSeconds prometheus.Histogram
|
||||
BatchResultBytes prometheus.Histogram
|
||||
handler http.Handler
|
||||
}
|
||||
|
||||
@@ -94,7 +101,40 @@ func New(src Source) *Set {
|
||||
Name: "sift_pdf_cache_hits_total",
|
||||
Help: "PDF answers served from the document cache.",
|
||||
})
|
||||
reg.MustRegister(requests, duration, pdfTotal, pdfSeconds, pdfPages, pdfWait, pdfCacheHits, &collector{src: src})
|
||||
batchJobs := prometheus.NewCounterVec(prometheus.CounterOpts{
|
||||
Name: "sift_batch_jobs_total",
|
||||
Help: "Batch jobs by outcome.",
|
||||
}, []string{"outcome"})
|
||||
batchTargets := prometheus.NewCounterVec(prometheus.CounterOpts{
|
||||
Name: "sift_batch_targets_total",
|
||||
Help: "Batch targets by outcome.",
|
||||
}, []string{"outcome"})
|
||||
batchQueueDepth := prometheus.NewGauge(prometheus.GaugeOpts{
|
||||
Name: "sift_batch_queue_depth",
|
||||
Help: "Live batch jobs (queued+running).",
|
||||
})
|
||||
batchTargetsInflight := prometheus.NewGauge(prometheus.GaugeOpts{
|
||||
Name: "sift_batch_targets_inflight",
|
||||
Help: "Batch targets currently being fetched.",
|
||||
})
|
||||
batchJobSeconds := prometheus.NewHistogram(prometheus.HistogramOpts{
|
||||
Name: "sift_batch_job_seconds",
|
||||
Help: "Batch wall time from submit to terminal status.",
|
||||
Buckets: durationBuckets,
|
||||
})
|
||||
batchTargetSeconds := prometheus.NewHistogram(prometheus.HistogramOpts{
|
||||
Name: "sift_batch_target_seconds",
|
||||
Help: "Batch target fetch wall time.",
|
||||
Buckets: durationBuckets,
|
||||
})
|
||||
batchResultBytes := prometheus.NewHistogram(prometheus.HistogramOpts{
|
||||
Name: "sift_batch_result_bytes",
|
||||
Help: "Stored bytes per batch target result.",
|
||||
Buckets: batchResultBytesBuckets,
|
||||
})
|
||||
reg.MustRegister(requests, duration, pdfTotal, pdfSeconds, pdfPages, pdfWait, pdfCacheHits,
|
||||
batchJobs, batchTargets, batchQueueDepth, batchTargetsInflight,
|
||||
batchJobSeconds, batchTargetSeconds, batchResultBytes, &collector{src: src})
|
||||
return &Set{
|
||||
Reg: reg,
|
||||
Requests: requests,
|
||||
@@ -104,6 +144,13 @@ func New(src Source) *Set {
|
||||
PDFPages: pdfPages,
|
||||
PDFWait: pdfWait,
|
||||
PDFCacheHits: pdfCacheHits,
|
||||
BatchJobsTotal: batchJobs,
|
||||
BatchTargetsTotal: batchTargets,
|
||||
BatchQueueDepth: batchQueueDepth,
|
||||
BatchTargetsInflight: batchTargetsInflight,
|
||||
BatchJobSeconds: batchJobSeconds,
|
||||
BatchTargetSeconds: batchTargetSeconds,
|
||||
BatchResultBytes: batchResultBytes,
|
||||
handler: promhttp.HandlerFor(reg, promhttp.HandlerOpts{EnableOpenMetrics: false}),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -234,11 +234,14 @@ func TestGrepLimits(t *testing.T) {
|
||||
if _, err := pdfdoc.Render(d, pdfdoc.Options{Grep: "(unclosed", OutCap: 1000}); !errors.Is(err, pdfdoc.ErrBadGrep) {
|
||||
t.Errorf("bad regexp: %v, want ErrBadGrep", err)
|
||||
}
|
||||
// RE2 is linear: the classic catastrophic pattern over a 1 MB page finishes quickly.
|
||||
// RE2 is linear: the classic catastrophic pattern over a 1 MB page finishes. A backtracking
|
||||
// engine would not finish at all on this input, so the bound only has to be finite and far
|
||||
// above normal: 10 s. (1 s was too tight under -race, which slows regexp ~15–20x: talos
|
||||
// measured 1.5–1.9 s, 2026-09-26.)
|
||||
evil := &pdftext.Doc{Meta: pdftext.Meta{Pages: 1}, Text: []string{strings.Repeat("a", 1<<20) + "!"}}
|
||||
start := time.Now()
|
||||
_, _ = pdfdoc.Render(evil, pdfdoc.Options{Grep: "(a+)+$", OutCap: 1000})
|
||||
if d := time.Since(start); d > time.Second {
|
||||
if d := time.Since(start); d > 10*time.Second {
|
||||
t.Errorf("(a+)+$ took %v", d)
|
||||
}
|
||||
// Output stays under the cap.
|
||||
|
||||
@@ -0,0 +1,166 @@
|
||||
package search
|
||||
|
||||
import (
|
||||
"net/url"
|
||||
"strings"
|
||||
)
|
||||
|
||||
func resultHost(raw string) string {
|
||||
u, err := url.Parse(raw)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
h := strings.ToLower(u.Hostname())
|
||||
return strings.TrimSuffix(h, ".")
|
||||
}
|
||||
|
||||
func hostMatch(host, d string) bool {
|
||||
return host == d || strings.HasSuffix(host, "."+d)
|
||||
}
|
||||
|
||||
func filterDomains(hits []Hit, f Filter) []Hit {
|
||||
if len(f.IncludeDomains) == 0 && len(f.ExcludeDomains) == 0 {
|
||||
return hits
|
||||
}
|
||||
mode := f.IncludeDomainsMode
|
||||
if mode == "" {
|
||||
mode = "restrict"
|
||||
}
|
||||
restrict := mode != "prefer" && len(f.IncludeDomains) > 0
|
||||
out := make([]Hit, 0, len(hits))
|
||||
for _, h := range hits {
|
||||
host := resultHost(h.URL)
|
||||
if host == "" {
|
||||
if restrict {
|
||||
continue
|
||||
}
|
||||
} else {
|
||||
drop := false
|
||||
for _, d := range f.ExcludeDomains {
|
||||
if hostMatch(host, d) {
|
||||
drop = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if drop {
|
||||
continue
|
||||
}
|
||||
if restrict {
|
||||
ok := false
|
||||
for _, d := range f.IncludeDomains {
|
||||
if hostMatch(host, d) {
|
||||
ok = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
}
|
||||
}
|
||||
out = append(out, h)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func filterMinScore(hits []Hit, min float64) []Hit {
|
||||
if min <= 0 {
|
||||
return hits
|
||||
}
|
||||
out := make([]Hit, 0, len(hits))
|
||||
for _, h := range hits {
|
||||
if h.Score >= min {
|
||||
out = append(out, h)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func dedupeStrings(in []string) []string {
|
||||
seen := map[string]bool{}
|
||||
out := make([]string, 0, len(in))
|
||||
for _, s := range in {
|
||||
if s == "" || seen[s] {
|
||||
continue
|
||||
}
|
||||
seen[s] = true
|
||||
out = append(out, s)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func searxQueryAndNotes(q string, f Filter) (string, []string) {
|
||||
var notes []string
|
||||
if f.IncludeDomainsMode == "prefer" && len(f.IncludeDomains) > 0 {
|
||||
notes = append(notes, "searxng: include_domains_mode=prefer not applied")
|
||||
} else if n := len(f.IncludeDomains); n == 1 {
|
||||
q = "site:" + f.IncludeDomains[0] + " " + q
|
||||
} else if n > 1 {
|
||||
max := n
|
||||
if max > searxSiteCap {
|
||||
max = searxSiteCap
|
||||
notes = append(notes, "searxng: extra include_domains dropped")
|
||||
}
|
||||
var b strings.Builder
|
||||
for i := 0; i < max; i++ {
|
||||
if i > 0 {
|
||||
b.WriteString(" OR ")
|
||||
}
|
||||
b.WriteString("site:")
|
||||
b.WriteString(f.IncludeDomains[i])
|
||||
}
|
||||
b.WriteByte(' ')
|
||||
b.WriteString(q)
|
||||
q = b.String()
|
||||
}
|
||||
if f.Topic == "finance" {
|
||||
notes = append(notes, "searxng: topic ignored")
|
||||
}
|
||||
if f.StartDate != "" || f.EndDate != "" {
|
||||
notes = append(notes, "searxng: start_date/end_date not applied")
|
||||
}
|
||||
if f.MinScore > 0 {
|
||||
notes = append(notes, "searxng: min_score not applied")
|
||||
}
|
||||
return q, notes
|
||||
}
|
||||
|
||||
func braveNotes(f Filter) []string {
|
||||
var notes []string
|
||||
if f.TimeRange != "" {
|
||||
notes = append(notes, "brave_llm: time_range not applied")
|
||||
}
|
||||
if f.Topic != "" && f.Topic != "general" {
|
||||
notes = append(notes, "brave_llm: topic not applied")
|
||||
}
|
||||
if f.ExactMatch {
|
||||
notes = append(notes, "brave_llm: exact_match not applied")
|
||||
}
|
||||
if f.IncludeDomainsMode == "prefer" && len(f.IncludeDomains) > 0 {
|
||||
notes = append(notes, "brave_llm: include_domains_mode=prefer not applied")
|
||||
}
|
||||
if f.MinScore > 0 {
|
||||
notes = append(notes, "brave_llm: min_score not applied")
|
||||
}
|
||||
return notes
|
||||
}
|
||||
|
||||
func FilterKinds(f Filter) string {
|
||||
var k []string
|
||||
if len(f.IncludeDomains) > 0 || len(f.ExcludeDomains) > 0 {
|
||||
k = append(k, "domains")
|
||||
}
|
||||
if f.TimeRange != "" || f.StartDate != "" || f.EndDate != "" {
|
||||
k = append(k, "time")
|
||||
}
|
||||
if f.Topic != "" && f.Topic != "general" {
|
||||
k = append(k, "topic")
|
||||
}
|
||||
if f.ExactMatch {
|
||||
k = append(k, "exact")
|
||||
}
|
||||
if f.MinScore > 0 {
|
||||
k = append(k, "score")
|
||||
}
|
||||
return strings.Join(k, ",")
|
||||
}
|
||||
@@ -2,7 +2,6 @@ package search
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"sort"
|
||||
"sync"
|
||||
@@ -90,104 +89,7 @@ func (e *Engine) timeout() time.Duration {
|
||||
}
|
||||
|
||||
func (e *Engine) Run(ctx context.Context, query string, k int, doRerank bool, caller, class string) RunResult {
|
||||
start := e.now()
|
||||
k = ClampK(k, e.DefaultK)
|
||||
if e.DefaultK <= 0 {
|
||||
k = ClampK(k, DefaultK)
|
||||
}
|
||||
if class == "" {
|
||||
class = dcache.ClassFleet
|
||||
}
|
||||
|
||||
attempted := e.attemptedNames(class)
|
||||
key := dcache.SearchKey(query, attempted, k, e.now(), class)
|
||||
if b, ok := e.Cache.GetSearch(key); ok {
|
||||
var hits []Hit
|
||||
if json.Unmarshal(b, &hits) == nil {
|
||||
e.mu.Lock()
|
||||
e.CacheHit++
|
||||
e.mu.Unlock()
|
||||
return RunResult{Hits: clampHits(hits, k), Cached: true, Attempted: attempted, SearchMS: e.now().Sub(start).Milliseconds()}
|
||||
}
|
||||
}
|
||||
e.mu.Lock()
|
||||
e.CacheMiss++
|
||||
e.mu.Unlock()
|
||||
|
||||
type one struct {
|
||||
name string
|
||||
out SourceOut
|
||||
skip bool
|
||||
}
|
||||
outs := make([]one, len(e.Sources))
|
||||
var wg sync.WaitGroup
|
||||
for i, src := range e.Sources {
|
||||
name := src.Name()
|
||||
if name == NameTavily && e.tavilyCapped(class) {
|
||||
outs[i] = one{name: name, out: SourceOut{Warn: "tavily: daily cap"}, skip: true}
|
||||
e.addCapped(NameTavily)
|
||||
continue
|
||||
}
|
||||
wg.Add(1)
|
||||
go func(i int, src Source, name string) {
|
||||
defer wg.Done()
|
||||
sctx, cancel := context.WithTimeout(ctx, e.timeout())
|
||||
defer cancel()
|
||||
outs[i] = one{name: name, out: src.Search(sctx, query, k)}
|
||||
}(i, src, name)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
var lists []namedHits
|
||||
var hits []Hit
|
||||
var warns []string
|
||||
var okNames, failNames, att []string
|
||||
degraded := false
|
||||
for _, o := range outs {
|
||||
if o.skip {
|
||||
if o.out.Warn != "" {
|
||||
warns = append(warns, o.out.Warn)
|
||||
}
|
||||
continue
|
||||
}
|
||||
att = append(att, o.name)
|
||||
e.addQuery(caller, o.name)
|
||||
if o.out.Warn != "" {
|
||||
warns = append(warns, o.out.Warn)
|
||||
failNames = append(failNames, o.name)
|
||||
degraded = true
|
||||
e.addErr(o.name, o.out.Warn)
|
||||
continue
|
||||
}
|
||||
okNames = append(okNames, o.name)
|
||||
lists = append(lists, namedHits{name: o.name, hits: o.out.Hits})
|
||||
if o.out.Credits > 0 {
|
||||
e.addCredits(caller, o.name, o.out.Credits)
|
||||
if o.name == NameTavily && e.Cache != nil {
|
||||
e.Cache.AddSearchCredits(tavilyCapSource(class), e.now().UTC().Format("2006-01-02"), o.out.Credits)
|
||||
}
|
||||
}
|
||||
}
|
||||
hits = mergeSearch(e, lists, doRerank)
|
||||
|
||||
var rerankMS int64
|
||||
hits = clampHits(hits, k)
|
||||
|
||||
if !degraded && e.Cache != nil && len(att) > 0 {
|
||||
payload, _ := json.Marshal(hits)
|
||||
ck := dcache.SearchKey(query, att, k, e.now(), class)
|
||||
e.Cache.PutSearch(ck, payload, e.ttl())
|
||||
}
|
||||
|
||||
return RunResult{
|
||||
Hits: hits,
|
||||
Warnings: warns,
|
||||
SourcesOK: okNames,
|
||||
SourcesFailed: failNames,
|
||||
SearchMS: e.now().Sub(start).Milliseconds(),
|
||||
RerankMS: rerankMS,
|
||||
Attempted: att,
|
||||
}
|
||||
return e.RunFilter(ctx, []string{query}, k, doRerank, caller, class, Filter{})
|
||||
}
|
||||
|
||||
func (e *Engine) attemptedNames(class string) []string {
|
||||
|
||||
@@ -14,12 +14,13 @@ type fakeSrc struct {
|
||||
name string
|
||||
hits []Hit
|
||||
warn string
|
||||
notes []string
|
||||
delay time.Duration
|
||||
credits int
|
||||
}
|
||||
|
||||
func (f fakeSrc) Name() string { return f.name }
|
||||
func (f fakeSrc) Search(ctx context.Context, query string, k int) SourceOut {
|
||||
func (f fakeSrc) Search(ctx context.Context, query string, k int, _ Filter) SourceOut {
|
||||
if f.delay > 0 {
|
||||
select {
|
||||
case <-time.After(f.delay):
|
||||
@@ -27,7 +28,7 @@ func (f fakeSrc) Search(ctx context.Context, query string, k int) SourceOut {
|
||||
return SourceOut{Warn: f.name + ": timeout"}
|
||||
}
|
||||
}
|
||||
return SourceOut{Hits: f.hits, Warn: f.warn, Credits: f.credits}
|
||||
return SourceOut{Hits: f.hits, Warn: f.warn, Credits: f.credits, Notes: f.notes}
|
||||
}
|
||||
|
||||
type stubEmb struct {
|
||||
|
||||
@@ -0,0 +1,256 @@
|
||||
package search
|
||||
|
||||
import (
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
const (
|
||||
maxDomainList = 20
|
||||
maxHostLen = 253
|
||||
searxSiteCap = 3
|
||||
)
|
||||
|
||||
var timeRangeLong = map[string]string{
|
||||
"day": "day", "d": "day",
|
||||
"week": "week", "w": "week",
|
||||
"month": "month", "m": "month",
|
||||
"year": "year", "y": "year",
|
||||
}
|
||||
|
||||
// NormalizeSearch checks query-vs-queries and filter fields. errMsg is
|
||||
// "invalid query" or "invalid search"; hint is empty or a field name.
|
||||
func NormalizeSearch(query string, queries []string, f Filter) (qs []string, out Filter, errMsg, hint string) {
|
||||
qs, errMsg = normalizeQueries(query, queries)
|
||||
if errMsg != "" {
|
||||
return nil, Filter{}, errMsg, ""
|
||||
}
|
||||
out, hint = normalizeFilter(f)
|
||||
if hint != "" {
|
||||
return nil, Filter{}, "invalid search", hint
|
||||
}
|
||||
if out.ExactMatch {
|
||||
for _, q := range qs {
|
||||
if !hasQuotedPhrase(q) {
|
||||
return nil, Filter{}, "invalid search", "exact_match"
|
||||
}
|
||||
}
|
||||
}
|
||||
return qs, out, "", ""
|
||||
}
|
||||
|
||||
func normalizeQueries(query string, queries []string) ([]string, string) {
|
||||
q := strings.TrimSpace(query)
|
||||
qs := make([]string, 0, len(queries))
|
||||
for _, s := range queries {
|
||||
qs = append(qs, strings.TrimSpace(s))
|
||||
}
|
||||
hasQ := q != ""
|
||||
hasQS := len(qs) > 0
|
||||
if hasQ && hasQS {
|
||||
return nil, "invalid query"
|
||||
}
|
||||
if !hasQ && !hasQS {
|
||||
return nil, "invalid query"
|
||||
}
|
||||
if hasQ {
|
||||
qs = []string{q}
|
||||
}
|
||||
if len(qs) < 1 || len(qs) > 4 {
|
||||
return nil, "invalid query"
|
||||
}
|
||||
for _, s := range qs {
|
||||
if s == "" || utf8.RuneCountInString(s) > MaxQueryRunes {
|
||||
return nil, "invalid query"
|
||||
}
|
||||
}
|
||||
return qs, ""
|
||||
}
|
||||
|
||||
func normalizeFilter(f Filter) (Filter, string) {
|
||||
inc, hint := normalizeDomains(f.IncludeDomains, "include_domains")
|
||||
if hint != "" {
|
||||
return Filter{}, hint
|
||||
}
|
||||
exc, hint := normalizeDomains(f.ExcludeDomains, "exclude_domains")
|
||||
if hint != "" {
|
||||
return Filter{}, hint
|
||||
}
|
||||
mode := strings.TrimSpace(f.IncludeDomainsMode)
|
||||
switch mode {
|
||||
case "", "restrict", "prefer":
|
||||
default:
|
||||
return Filter{}, "include_domains_mode"
|
||||
}
|
||||
if mode != "" && len(inc) == 0 {
|
||||
return Filter{}, "include_domains_mode"
|
||||
}
|
||||
if len(inc) > 0 && mode == "" {
|
||||
mode = "restrict"
|
||||
}
|
||||
|
||||
tr := strings.TrimSpace(f.TimeRange)
|
||||
if tr != "" {
|
||||
long, ok := timeRangeLong[tr]
|
||||
if !ok {
|
||||
return Filter{}, "time_range"
|
||||
}
|
||||
tr = long
|
||||
}
|
||||
start := strings.TrimSpace(f.StartDate)
|
||||
end := strings.TrimSpace(f.EndDate)
|
||||
if start != "" {
|
||||
if !validDay(start) {
|
||||
return Filter{}, "start_date"
|
||||
}
|
||||
}
|
||||
if end != "" {
|
||||
if !validDay(end) {
|
||||
return Filter{}, "end_date"
|
||||
}
|
||||
}
|
||||
if (start != "" || end != "") && tr != "" {
|
||||
return Filter{}, "time_range"
|
||||
}
|
||||
if start != "" && end != "" {
|
||||
st, _ := time.Parse("2006-01-02", start)
|
||||
en, _ := time.Parse("2006-01-02", end)
|
||||
if en.Before(st) {
|
||||
return Filter{}, "end_date"
|
||||
}
|
||||
}
|
||||
|
||||
topic := strings.TrimSpace(f.Topic)
|
||||
switch topic {
|
||||
case "", "general", "news", "finance":
|
||||
default:
|
||||
return Filter{}, "topic"
|
||||
}
|
||||
|
||||
ms := f.MinScore
|
||||
if ms < 0 || ms > 1 {
|
||||
return Filter{}, "min_score"
|
||||
}
|
||||
|
||||
return Filter{
|
||||
IncludeDomains: inc,
|
||||
ExcludeDomains: exc,
|
||||
IncludeDomainsMode: mode,
|
||||
TimeRange: tr,
|
||||
StartDate: start,
|
||||
EndDate: end,
|
||||
Topic: topic,
|
||||
ExactMatch: f.ExactMatch,
|
||||
MinScore: ms,
|
||||
}, ""
|
||||
}
|
||||
|
||||
func normalizeDomains(in []string, hint string) ([]string, string) {
|
||||
if len(in) == 0 {
|
||||
return nil, ""
|
||||
}
|
||||
out := make([]string, 0, len(in))
|
||||
seen := map[string]bool{}
|
||||
for _, raw := range in {
|
||||
s := strings.TrimSpace(raw)
|
||||
if s == "" {
|
||||
continue
|
||||
}
|
||||
d, ok := normalizeDomain(s)
|
||||
if !ok {
|
||||
return nil, hint
|
||||
}
|
||||
if seen[d] {
|
||||
continue
|
||||
}
|
||||
seen[d] = true
|
||||
out = append(out, d)
|
||||
}
|
||||
if len(out) == 0 {
|
||||
return nil, ""
|
||||
}
|
||||
sort.Strings(out)
|
||||
if len(out) > maxDomainList {
|
||||
return nil, hint
|
||||
}
|
||||
return out, ""
|
||||
}
|
||||
|
||||
func normalizeDomain(s string) (string, bool) {
|
||||
s = strings.ToLower(strings.TrimSpace(s))
|
||||
s = strings.TrimSuffix(s, ".")
|
||||
if s == "" || len(s) > maxHostLen {
|
||||
return "", false
|
||||
}
|
||||
if strings.ContainsAny(s, " :/@*?#[]") {
|
||||
return "", false
|
||||
}
|
||||
labels := strings.Split(s, ".")
|
||||
if len(labels) == 0 {
|
||||
return "", false
|
||||
}
|
||||
for _, lab := range labels {
|
||||
if !validLabel(lab) {
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
return s, true
|
||||
}
|
||||
|
||||
func validLabel(lab string) bool {
|
||||
if lab == "" || len(lab) > 63 {
|
||||
return false
|
||||
}
|
||||
if lab[0] == '-' || lab[len(lab)-1] == '-' {
|
||||
return false
|
||||
}
|
||||
for i := 0; i < len(lab); i++ {
|
||||
c := lab[i]
|
||||
if c >= 'a' && c <= 'z' || c >= '0' && c <= '9' || c == '-' {
|
||||
continue
|
||||
}
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func validDay(s string) bool {
|
||||
t, err := time.Parse("2006-01-02", s)
|
||||
return err == nil && t.Format("2006-01-02") == s
|
||||
}
|
||||
|
||||
// hasQuotedPhrase reports a balanced ASCII "…" pair with ≥1 rune inside and no unpaired quotes.
|
||||
func hasQuotedPhrase(q string) bool {
|
||||
in := false
|
||||
n := 0
|
||||
pairs := 0
|
||||
for _, r := range q {
|
||||
if r != '"' {
|
||||
if in {
|
||||
n++
|
||||
}
|
||||
continue
|
||||
}
|
||||
if !in {
|
||||
in = true
|
||||
n = 0
|
||||
continue
|
||||
}
|
||||
if n == 0 {
|
||||
return false
|
||||
}
|
||||
in = false
|
||||
pairs++
|
||||
}
|
||||
return !in && pairs >= 1
|
||||
}
|
||||
|
||||
func QueryRuneLen(qs []string) int {
|
||||
n := 0
|
||||
for _, q := range qs {
|
||||
n += utf8.RuneCountInString(q)
|
||||
}
|
||||
return n
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user