Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions core/middleware.go
Original file line number Diff line number Diff line change
Expand Up @@ -65,10 +65,10 @@ func RequestContextMiddleware() fiber.Handler {

// RequestTimeoutMiddleware bounds wall-clock time per request by attaching a
// deadline to the user context, which fasthttp never cancels on client
// disconnect. /mega/* (MegaTimeout) and /extract (batch budget) are exempt.
// disconnect. /mega/* (MegaTimeout) and /extract* (batch budget) are exempt.
func RequestTimeoutMiddleware(timeout time.Duration) fiber.Handler {
return func(c *fiber.Ctx) error {
if strings.HasPrefix(c.Path(), "/mega/") || c.Path() == "/extract" {
if strings.HasPrefix(c.Path(), "/mega/") || c.Path() == "/extract" || c.Path() == "/extract/batch" {
return c.Next()
}
ctx, cancel := context.WithTimeout(c.UserContext(), timeout)
Expand Down
1 change: 1 addition & 0 deletions core/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -235,6 +235,7 @@ func NewServerWithOptions(host string, port int, opts ServerOptions, searchEngin
serv.app.Get("/mega/engines", serv.handleListEngines)
serv.app.Get("/extract", serv.handleExtract)
serv.app.Post("/extract", serv.handleExtract)
serv.app.Post("/extract/batch", serv.handleBatchExtract)

return &serv
}
Expand Down
126 changes: 126 additions & 0 deletions core/server_extract.go
Original file line number Diff line number Diff line change
Expand Up @@ -418,3 +418,129 @@ func SanitizeExtractError(err error) string {
}
return msg
}

const maxBatchExtractURLs = 20

type batchExtractPayload struct {
URLs []string `json:"urls"`
Mode string `json:"mode"`
}

// batchExtractItem is a response item for a single extracted URL, using
// page_content/metadata keys
type batchExtractItem struct {
PageContent string `json:"page_content"`
Metadata map[string]string `json:"metadata"`
}

func (s *Server) handleBatchExtract(c *fiber.Ctx) error {
requestCtx := withRequestUsage(c.UserContext(), "extract-batch")
c.SetUserContext(requestCtx)
defer setNetworkBytesHeader(c, requestCtx)
defer setBrowserProfileHeader(c, requestCtx)

cfg := s.opts.Extract.Normalized()
if !cfg.Enabled {
return &APIError{HTTPStatus: fiber.StatusNotFound, ErrorCode: "not_found", Message: "Extraction is disabled"}
}

var body batchExtractPayload
if len(c.Body()) == 0 {
return errInvalidParam("request body is required")
}
if err := c.BodyParser(&body); err != nil {
return errInvalidParam("invalid JSON body")
}

// Deduplicate and normalize URLs.
seen := make(map[string]struct{}, len(body.URLs))
var urls []string
for _, raw := range body.URLs {
u := extractpkg.NormalizeURL(strings.TrimSpace(raw))
if u == "" {
continue
}
if _, dup := seen[u]; dup {
continue
}
seen[u] = struct{}{}
urls = append(urls, u)
}
if len(urls) == 0 {
return errInvalidParam("urls array is required and must contain at least one valid URL")
}
if len(urls) > maxBatchExtractURLs {
return errInvalidParam(fmt.Sprintf("urls array exceeds maximum of %d", maxBatchExtractURLs))
}

// Validate all URLs upfront.
for _, u := range urls {
if err := validateExtractTargetURL(c.UserContext(), u, cfg.AllowPrivateNetworks); err != nil {
return &APIError{HTTPStatus: fiber.StatusBadRequest, ErrorCode: "invalid_extract_url", Message: err.Error()}
}
}

mode := firstNonEmpty(body.Mode, cfg.DefaultMode)
extractor := s.newExtractor()
results := make([]batchExtractItem, len(urls))

// Concurrent extraction with bounded parallelism (same pattern as
// EnrichEnvelopeWithExtraction).
ctx, cancel := context.WithTimeout(c.UserContext(), cfg.BatchTimeout(len(urls)))
Comment thread
AIWintermuteAI marked this conversation as resolved.
defer cancel()

sem := make(chan struct{}, cfg.MaxConcurrent)
var wg sync.WaitGroup
for i, u := range urls {
wg.Add(1)
sem <- struct{}{}
go func(idx int, url string) {
defer wg.Done()
defer func() { <-sem }()

if err := ctx.Err(); err != nil {
results[idx] = batchExtractItem{
PageContent: "",
Metadata: map[string]string{"source": url, "error": "batch timeout"},
}
return
}

req := extractpkg.ExtractRequest{
URL: url,
Mode: extractpkg.Mode(mode),
ProxyURL: "",
LangCode: strings.TrimSpace(c.Query("lang")),
Comment thread
AIWintermuteAI marked this conversation as resolved.
Timeout: cfg.Timeout,
MaxBytes: cfg.MaxBytes,
}
result, err := extractor.Extract(ctx, req)
if err != nil {
results[idx] = batchExtractItem{
PageContent: "",
Metadata: map[string]string{
"source": url,
"error": SanitizeExtractError(err),
},
}
return
}
results[idx] = batchExtractItem{
PageContent: result.Markdown,
Metadata: map[string]string{
"source": url,
"title": result.Title,
"description": result.Description,
"lang": result.Lang,
"canonical": result.Canonical,
"mode_used": result.Meta.ModeUsed,
"fetched_at": result.Meta.FetchedAt,
"took_ms": fmt.Sprintf("%d", result.Meta.TookMs),
},
}
}(i, u)
}
wg.Wait()

return c.JSON(results)
}
148 changes: 148 additions & 0 deletions core/server_extract_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,11 @@ package core

import (
"context"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"reflect"
"strings"
"testing"
"time"
Expand Down Expand Up @@ -109,3 +112,148 @@ func TestValidateExtractTargetURLNormalizesBarePublicIP(t *testing.T) {
t.Fatalf("expected bare public IP target to validate after scheme normalization: %v", err)
}
}

func TestBatchExtractSingleURL(t *testing.T) {
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`<html><body><article><h1>Test Page</h1><p>This is a test page with enough content to pass the minimum runes threshold for extraction in batch mode.</p></article></body></html>`))
}))
defer target.Close()

opts := DefaultServerOptions()
opts.Extract = extractpkg.Config{
Enabled: true,
DefaultMode: string(extractpkg.ModeFast),
Timeout: time.Second,
MaxBytes: 256 * 1024,
MaxConcurrent: 2,
AllowPrivateNetworks: true,
}
s := NewServerWithOptions("127.0.0.1", 0, opts)

body := fmt.Sprintf(`{"urls":["%s"]}`, target.URL)
req, err := http.NewRequest(http.MethodPost, "/extract/batch", strings.NewReader(body))
if err != nil {
t.Fatal(err)
}
req.Header.Set("Content-Type", "application/json")

resp, err := s.app.Test(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()

if resp.StatusCode != http.StatusOK {
t.Fatalf("status = %d, want %d", resp.StatusCode, http.StatusOK)
}

var results []map[string]interface{}
if err := json.NewDecoder(resp.Body).Decode(&results); err != nil {
t.Fatalf("decode response: %v", err)
}
if len(results) != 1 {
t.Fatalf("results count = %d, want 1", len(results))
}
if _, ok := results[0]["page_content"]; !ok {
t.Fatalf("expected page_content key, got keys: %v", reflect.ValueOf(results[0]).MapKeys())
}
if _, ok := results[0]["metadata"]; !ok {
t.Fatalf("expected metadata key, got keys: %v", reflect.ValueOf(results[0]).MapKeys())
}
}

func TestBatchExtractHandlesMultipleURLs(t *testing.T) {
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
_, _ = w.Write([]byte(`<html><body><article><h1>Multi</h1><p>Page with sufficient content for batch extraction test that verifies concurrent processing works correctly.</p></article></body></html>`))
}))
defer target.Close()

opts := DefaultServerOptions()
opts.Extract = extractpkg.Config{
Enabled: true,
DefaultMode: string(extractpkg.ModeFast),
Timeout: time.Second,
MaxBytes: 256 * 1024,
MaxConcurrent: 2,
AllowPrivateNetworks: true,
}
s := NewServerWithOptions("127.0.0.1", 0, opts)

body := fmt.Sprintf(`{"urls":["%s/1","%s/2","%s/3"]}`, target.URL, target.URL, target.URL)
req, err := http.NewRequest(http.MethodPost, "/extract/batch", strings.NewReader(body))
if err != nil {
t.Fatal(err)
}
req.Header.Set("Content-Type", "application/json")

resp, err := s.app.Test(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()

if resp.StatusCode != http.StatusOK {
t.Fatalf("status = %d, want %d", resp.StatusCode, http.StatusOK)
}

var results []map[string]interface{}
if err := json.NewDecoder(resp.Body).Decode(&results); err != nil {
t.Fatalf("decode response: %v", err)
}
if len(results) != 3 {
t.Fatalf("results count = %d, want 3", len(results))
}
}

func TestBatchExtractRejectsEmptyURLs(t *testing.T) {
opts := DefaultServerOptions()
opts.Extract = extractpkg.DefaultConfig()
s := NewServerWithOptions("127.0.0.1", 0, opts)

req, err := http.NewRequest(http.MethodPost, "/extract/batch", strings.NewReader(`{"urls":[]}`))
if err != nil {
t.Fatal(err)
}
req.Header.Set("Content-Type", "application/json")

resp, err := s.app.Test(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()

if resp.StatusCode != http.StatusBadRequest {
t.Fatalf("status = %d, want %d", resp.StatusCode, http.StatusBadRequest)
}
}

func TestBatchExtractRejectsURLsOverLimit(t *testing.T) {
opts := DefaultServerOptions()
opts.Extract = extractpkg.DefaultConfig()
s := NewServerWithOptions("127.0.0.1", 0, opts)

// Build 21 URLs (limit is 20)
urls := make([]string, 21)
for i := 0; i < 21; i++ {
urls[i] = fmt.Sprintf("https://example.com/%d", i)
}
body, _ := json.Marshal(map[string][]string{"urls": urls})

req, err := http.NewRequest(http.MethodPost, "/extract/batch", strings.NewReader(string(body)))
if err != nil {
t.Fatal(err)
}
req.Header.Set("Content-Type", "application/json")

resp, err := s.app.Test(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()

if resp.StatusCode != http.StatusBadRequest {
t.Fatalf("status = %d, want %d", resp.StatusCode, http.StatusBadRequest)
}
}
7 changes: 5 additions & 2 deletions core/server_timeout_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,8 @@ import (
)

// FP-2: every endpoint that doesn't manage its own deadline budget must get
// one from RequestTimeoutMiddleware; /mega/* (MegaTimeout) and /extract
// (batch budget) keep theirs.
// one from RequestTimeoutMiddleware; /mega/* (MegaTimeout) and /extract,
// /extract/batch (batch budget) keep theirs.
func TestRequestTimeoutMiddlewareSetsDeadlineExceptBudgetedPaths(t *testing.T) {
app := fiber.New()
app.Use(RequestTimeoutMiddleware(time.Minute))
Expand All @@ -32,12 +32,14 @@ func TestRequestTimeoutMiddlewareSetsDeadlineExceptBudgetedPaths(t *testing.T) {
app.Post("/google/parse", record("/google/parse"))
app.Get("/mega/search", record("/mega/search"))
app.Get("/extract", record("/extract"))
app.Post("/extract/batch", record("/extract/batch"))

for path, method := range map[string]string{
"/google/search": http.MethodGet,
"/google/parse": http.MethodPost,
"/mega/search": http.MethodGet,
"/extract": http.MethodGet,
"/extract/batch": http.MethodPost,
} {
req := httptest.NewRequest(method, path, nil)
resp, err := app.Test(req, -1)
Expand All @@ -54,6 +56,7 @@ func TestRequestTimeoutMiddlewareSetsDeadlineExceptBudgetedPaths(t *testing.T) {
"/google/parse": true,
"/mega/search": false,
"/extract": false,
"/extract/batch": false,
} {
if deadlines[path] != want {
t.Errorf("%s: deadline attached = %v, want %v", path, deadlines[path], want)
Expand Down
Loading
Loading