1use std::collections::HashMap;
5use std::future::Future;
6use std::pin::{Pin, pin};
7use std::sync::Arc;
8use std::sync::atomic::{AtomicBool, Ordering};
9use std::task::{Context, Poll};
10use std::time::{Duration, Instant};
11
12use axum::body::{Body, Bytes, HttpBody};
13use axum::extract::rejection::JsonRejection;
14use axum::extract::{Request, State};
15use axum::http::{HeaderMap, Method, StatusCode, header};
16use axum::middleware::{self, Next};
17use axum::response::{IntoResponse, Response};
18use axum::routing::get;
19use axum::{Json, Router};
20use http_body::{Frame, SizeHint};
21use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
22use hyper_util::server::conn::auto;
23use hyper_util::server::graceful::GracefulShutdown;
24use hyper_util::service::TowerToHyperService;
25use pagesift_store::Store;
26use tokio::net::TcpListener;
27use tokio::sync::Mutex;
28use tokio::time::Sleep;
29
30use crate::engine::Engine;
31use crate::error::{ApiError, ErrorCode};
32use crate::jobs::Jobs;
33use crate::limits::{Caller, KeyLimits, Limiter};
34use crate::metrics::Metrics;
35use crate::webhook::{WebhookConfig, WebhookSettings, Webhooks};
36
37pub const PAGE_SIZE: usize = 100;
39pub const DEFAULT_BATCH_CONCURRENCY: usize = 8;
43
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub struct Timeouts {
47 pub header_read: Duration,
49 pub body: Duration,
51 pub upload_body: Duration,
54 pub request: Duration,
57}
58
59impl Default for Timeouts {
60 fn default() -> Self {
61 Self {
62 header_read: Duration::from_secs(30),
63 body: Duration::from_secs(30),
64 upload_body: Duration::from_secs(120),
65 request: Duration::from_millis(crate::validate::MAX_SCRAPE_TIMEOUT_MS)
66 + Duration::from_secs(60),
67 }
68 }
69}
70
71#[derive(Clone, Default)]
72pub struct Config {
73 pub api_keys: Vec<String>,
75 pub public_url: Option<String>,
77 pub max_batch_concurrency: Option<usize>,
78 pub store: Option<Store>,
80 pub key_limits: HashMap<String, KeyLimits>,
82 pub default_limits: KeyLimits,
84 pub webhook: WebhookSettings,
85 pub timeouts: Timeouts,
87 pub ledger: Option<Arc<pagesift_cluster::Ledger>>,
90 pub job_mirror: Option<Arc<crate::JobMirror>>,
92 pub extra_metrics: Option<ExtraMetrics>,
94 pub monitor_leader: Option<crate::monitor::Leader>,
96}
97
98pub type ExtraMetrics = Arc<dyn Fn() -> String + Send + Sync>;
100
101impl std::fmt::Debug for Config {
102 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
103 f.debug_struct("Config")
104 .field("api_keys", &self.api_keys.len())
105 .field("public_url", &self.public_url)
106 .field("max_batch_concurrency", &self.max_batch_concurrency)
107 .field("store", &self.store.is_some())
108 .field("key_limits", &self.key_limits.len())
109 .field("default_limits", &self.default_limits)
110 .field("webhook", &self.webhook)
111 .field("timeouts", &self.timeouts)
112 .field("ledger", &self.ledger.is_some())
113 .field("job_mirror", &self.job_mirror.is_some())
114 .field("extra_metrics", &self.extra_metrics.is_some())
115 .finish()
116 }
117}
118
119#[derive(Clone)]
120pub struct AppState {
121 pub engine: Arc<dyn Engine>,
122 pub keys: Arc<HashMap<String, Arc<KeyLimits>>>,
123 pub default_limits: Arc<KeyLimits>,
124 pub public_url: Option<String>,
125 pub batch_concurrency: Option<usize>,
127 pub jobs: Jobs,
128 pub limiter: Arc<Limiter>,
129 pub webhooks: Webhooks,
130 pub extracts: Arc<Mutex<HashMap<String, ExtractJob>>>,
131 pub agents: crate::agent::Agents,
132 pub(crate) monitors: crate::monitor::Monitors,
133 pub metrics: Arc<Metrics>,
134 pub extra_metrics: Option<ExtraMetrics>,
135}
136
137impl AppState {
138 pub fn batch_concurrency_for(&self, requested: Option<usize>) -> usize {
142 batch_concurrency(
143 requested,
144 self.engine.admitted_concurrency(),
145 self.batch_concurrency,
146 )
147 }
148}
149
150fn batch_concurrency(
151 requested: Option<usize>,
152 admitted: Option<usize>,
153 cap: Option<usize>,
154) -> usize {
155 let admitted = admitted.unwrap_or(cap.unwrap_or(DEFAULT_BATCH_CONCURRENCY));
156 requested
157 .unwrap_or(admitted)
158 .min(admitted)
159 .min(cap.unwrap_or(usize::MAX))
160 .max(1)
161}
162
163pub struct ExtractJob {
164 pub status: &'static str,
165 pub(crate) owner: String,
166 pub data: Option<serde_json::Value>,
167 pub error: Option<ApiError>,
168 pub created: std::time::SystemTime,
169}
170
171pub fn router(engine: Arc<dyn Engine>, config: Config) -> Router {
173 let timeouts = config.timeouts;
174 let default_limits = Arc::new(config.default_limits);
175 let mut keys: HashMap<String, Arc<KeyLimits>> = config
176 .api_keys
177 .into_iter()
178 .map(|k| (k, default_limits.clone()))
179 .collect();
180 keys.extend(config.key_limits.into_iter().map(|(k, l)| (k, Arc::new(l))));
181
182 let accounts: HashMap<String, Arc<KeyLimits>> = keys
183 .iter()
184 .map(|(k, l)| {
185 (
186 Caller::new(Some(k), l.clone()).account.to_string(),
187 l.clone(),
188 )
189 })
190 .collect();
191
192 let metrics = Arc::new(Metrics::default());
193 let limiter = Arc::new(match config.ledger {
194 Some(ledger) => {
195 let mut all = accounts.clone();
197 all.entry("anonymous".into())
198 .or_insert_with(|| default_limits.clone());
199 Limiter::with_ledger(ledger, &all)
200 }
201 None => Limiter::new(config.store.clone()),
202 });
203 if let Some(s) = &config.store
205 && tokio::runtime::Handle::try_current().is_ok()
206 {
207 crate::attachments::sweep_expired(crate::attachments::Disk {
208 store: s.clone(),
209 mirror: config.job_mirror.clone(),
210 });
211 crate::attachments::backfill_access(crate::attachments::Disk {
212 store: s.clone(),
213 mirror: config.job_mirror.clone(),
214 });
215 }
216 let mut webhooks = Webhooks::new(config.webhook, metrics.clone());
217 webhooks.mirror = config.job_mirror.clone();
218 let agents = crate::agent::Agents::new(
219 config.store.clone(),
220 Some(webhooks.clone()),
221 Some(limiter.clone()),
222 );
223 let jobs = Jobs::new(
224 config.store,
225 limiter.clone(),
226 Some(webhooks.clone()),
227 metrics.clone(),
228 config.job_mirror,
229 );
230
231 if tokio::runtime::Handle::try_current().is_ok() {
232 let fallback = default_limits.clone();
233 jobs.restore(&engine, &|account| {
234 accounts
235 .get(account)
236 .cloned()
237 .unwrap_or_else(|| fallback.clone())
238 });
239 } else {
240 tracing::warn!("router built outside a tokio runtime: saved jobs are not resumed");
241 }
242
243 let monitors = crate::monitor::Monitors {
244 accounts: Arc::new(accounts),
245 default_limits: default_limits.clone(),
246 ..Default::default()
247 };
248 let state = AppState {
249 engine: engine.clone(),
250 keys: Arc::new(keys),
251 default_limits,
252 public_url: config
253 .public_url
254 .map(|u| u.trim_end_matches('/').to_string()),
255 batch_concurrency: config.max_batch_concurrency,
256 jobs,
257 limiter,
258 webhooks,
259 extracts: Arc::default(),
260 agents,
261 monitors,
262 metrics,
263 extra_metrics: config.extra_metrics,
264 };
265
266 if tokio::runtime::Handle::try_current().is_ok() {
268 crate::monitor::start_scheduler(state.clone(), config.monitor_leader.clone());
269 }
270 let api = crate::firecrawl::firecrawl_routes()
271 .merge(crate::jina::jina_routes())
272 .merge(crate::monitor::routes())
273 .route_layer(middleware::from_fn_with_state(state.clone(), auth));
274
275 Router::new()
276 .route("/healthz", get(|| async { "ok" }))
277 .route("/metrics", get(render_metrics))
278 .merge(api)
279 .with_state(state.clone())
280 .fallback(|| async { ApiError::new(ErrorCode::NotFound, "no such endpoint") })
281 .layer(middleware::from_fn(uniform_errors))
282 .layer(middleware::from_fn_with_state(timeouts, deadline))
283 .layer(middleware::from_fn_with_state(state.clone(), observe))
284}
285
286pub async fn serve(
288 listener: TcpListener,
289 app: Router,
290 shutdown: impl Future<Output = ()> + Send,
291) -> std::io::Result<()> {
292 serve_with(listener, app, Timeouts::default(), shutdown).await
293}
294
295const ACCEPT_RETRY_CAP: Duration = Duration::from_secs(1);
299
300pub async fn serve_with(
303 listener: TcpListener,
304 app: Router,
305 timeouts: Timeouts,
306 shutdown: impl Future<Output = ()> + Send,
307) -> std::io::Result<()> {
308 let mut builder = auto::Builder::new(TokioExecutor::new());
309 builder
310 .http1()
311 .timer(TokioTimer::new())
312 .header_read_timeout(timeouts.header_read);
313 let graceful = GracefulShutdown::new();
314 let mut shutdown = pin!(shutdown);
315 let released = Arc::new(tokio::sync::Notify::new());
317 loop {
318 let stream = tokio::select! {
319 r = listener.accept() => match r {
320 Ok((stream, _)) => stream,
321 Err(e) => {
322 if !matches!(
327 e.kind(),
328 std::io::ErrorKind::ConnectionRefused
329 | std::io::ErrorKind::ConnectionAborted
330 | std::io::ErrorKind::ConnectionReset
331 ) {
332 tracing::error!(error = %e, "accept failed");
333 tokio::select! {
334 () = released.notified() => {}
335 () = tokio::time::sleep(ACCEPT_RETRY_CAP) => {}
336 () = &mut shutdown => break,
337 }
338 }
339 continue;
340 }
341 },
342 () = &mut shutdown => break,
343 };
344 let _ = stream.set_nodelay(true);
346 let conn = builder
347 .serve_connection_with_upgrades(
348 TokioIo::new(stream),
349 TowerToHyperService::new(app.clone()),
350 )
351 .into_owned();
352 let conn = graceful.watch(conn);
353 let released = released.clone();
354 tokio::spawn(async move {
355 if let Err(e) = conn.await {
356 tracing::trace!(error = %e, "connection ended with an error");
357 }
358 released.notify_one();
359 });
360 }
361 drop(listener);
362 graceful.shutdown().await;
363 Ok(())
364}
365
366async fn deadline(State(t): State<Timeouts>, req: Request, next: Next) -> Response {
368 let fired = Arc::new(AtomicBool::new(false));
369 let (body_limit, request_limit) = if req.uri().path() == "/v2/parse" {
371 (
372 t.upload_body,
373 t.request + t.upload_body.saturating_sub(t.body),
374 )
375 } else {
376 (t.body, t.request)
377 };
378 let req = if req.body().is_end_stream() {
379 req
380 } else {
381 let fired = fired.clone();
382 req.map(|inner| {
383 Body::new(DeadlineBody {
384 inner,
385 sleep: Box::pin(tokio::time::sleep(body_limit)),
386 fired,
387 })
388 })
389 };
390 match tokio::time::timeout(request_limit, next.run(req)).await {
391 Ok(_) if fired.load(Ordering::Relaxed) => ApiError::new(
393 ErrorCode::Timeout,
394 format!(
395 "request body was not received within {} s",
396 body_limit.as_secs_f64()
397 ),
398 )
399 .into_response(),
400 Ok(resp) => resp,
401 Err(_) => ApiError::new(
402 ErrorCode::Timeout,
403 format!(
404 "request did not complete within {} s",
405 request_limit.as_secs_f64()
406 ),
407 )
408 .into_response(),
409 }
410}
411
412struct DeadlineBody {
414 inner: Body,
415 sleep: Pin<Box<Sleep>>,
416 fired: Arc<AtomicBool>,
417}
418
419impl HttpBody for DeadlineBody {
420 type Data = Bytes;
421 type Error = axum::Error;
422
423 fn poll_frame(
424 self: Pin<&mut Self>,
425 cx: &mut Context<'_>,
426 ) -> Poll<Option<Result<Frame<Bytes>, axum::Error>>> {
427 let this = self.get_mut();
428 if let Poll::Ready(frame) = Pin::new(&mut this.inner).poll_frame(cx) {
429 return Poll::Ready(frame);
430 }
431 if this.sleep.as_mut().poll(cx).is_ready() {
432 this.fired.store(true, Ordering::Relaxed);
433 return Poll::Ready(Some(Err(axum::Error::new("request body timed out"))));
434 }
435 Poll::Pending
436 }
437
438 fn is_end_stream(&self) -> bool {
439 self.inner.is_end_stream()
440 }
441
442 fn size_hint(&self) -> SizeHint {
443 self.inner.size_hint()
444 }
445}
446
447async fn uniform_errors(req: Request, next: Next) -> Response {
450 let resp = next.run(req).await;
451 let status = resp.status();
452 if !status.is_client_error() || resp.extensions().get::<ErrorCode>().is_some() {
454 return resp;
455 }
456 let plain = resp
457 .headers()
458 .get(header::CONTENT_TYPE)
459 .is_none_or(|v| v.as_bytes().starts_with(b"text/plain"));
460 if !plain {
461 return resp;
462 }
463 let (parts, body) = resp.into_parts();
464 let text = axum::body::to_bytes(body, 64 * 1024)
465 .await
466 .map(|b| String::from_utf8_lossy(&b).trim().to_string())
467 .unwrap_or_default();
468 let code = match status {
469 StatusCode::NOT_FOUND => ErrorCode::NotFound,
470 StatusCode::PAYLOAD_TOO_LARGE => ErrorCode::TooLarge,
471 StatusCode::UNAUTHORIZED => ErrorCode::Unauthorized,
472 _ => ErrorCode::InvalidRequest,
473 };
474 let message = if text.is_empty() {
475 status
476 .canonical_reason()
477 .unwrap_or("bad request")
478 .to_lowercase()
479 } else {
480 text
481 };
482 let mut out = ApiError::new(code, message).into_response();
483 *out.status_mut() = status;
484 for (name, value) in &parts.headers {
485 if name != header::CONTENT_TYPE && name != header::CONTENT_LENGTH {
486 out.headers_mut().append(name, value.clone());
487 }
488 }
489 out
490}
491
492async fn auth(
494 State(s): State<AppState>,
495 headers: HeaderMap,
496 mut req: Request,
497 next: Next,
498) -> Response {
499 let caller = if s.keys.is_empty() {
500 Caller::new(None, s.default_limits.clone())
501 } else {
502 let key = headers
503 .get(header::AUTHORIZATION)
504 .and_then(|v| v.to_str().ok())
505 .and_then(|v| v.strip_prefix("Bearer "))
506 .map(str::trim);
507 match key.and_then(|k| s.keys.get(k).map(|l| (k, l))) {
508 Some((k, l)) => Caller::new(Some(k), l.clone()),
509 None => {
510 return ApiError::new(ErrorCode::Unauthorized, "missing or invalid API key")
511 .into_response();
512 }
513 }
514 };
515 let path = req.uri().path();
517 if (req.method() == Method::POST || path.starts_with("/r/") || path.starts_with("/s/"))
518 && let Err(e) = s.limiter.check_rate(&caller)
519 {
520 let mut resp = e.into_response();
521 resp.extensions_mut().insert(caller);
522 return resp;
523 }
524 req.extensions_mut().insert(caller.clone());
525 let mut resp = next.run(req).await;
526 resp.extensions_mut().insert(caller);
527 resp
528}
529
530async fn observe(State(s): State<AppState>, req: Request, next: Next) -> Response {
532 let route = req
533 .extensions()
534 .get::<axum::extract::MatchedPath>()
535 .map(|p| p.as_str().to_string())
536 .unwrap_or_else(|| "unmatched".into());
537 let method = req.method().clone();
538 let started = Instant::now();
539 s.metrics.request_started();
540 let resp = next.run(req).await;
541 let elapsed = started.elapsed();
542 let status = resp.status().as_u16();
543 s.metrics.request_finished(&route, status, elapsed);
544 let code = resp.extensions().get::<ErrorCode>().map(|c| c.as_str());
545 if let Some(code) = code {
546 s.metrics.error(code, "http");
547 }
548 let account = resp.extensions().get::<Caller>().map(|c| c.account.clone());
549 let latency_ms = elapsed.as_secs_f64() * 1000.0;
550 if route.starts_with("/v2/") || route.starts_with("/v1/") || route.starts_with("/r") {
551 tracing::info!(
552 target: "pagesift::http",
553 %method, %route, status, latency_ms, account = account.as_deref(), code,
554 "request"
555 );
556 } else {
557 tracing::debug!(target: "pagesift::http", %method, %route, status, latency_ms, "request");
558 }
559 resp
560}
561
562pub fn body<T>(r: Result<Json<T>, JsonRejection>) -> Result<T, ApiError> {
563 r.map(|Json(v)| v)
564 .map_err(|e| ApiError::invalid(format!("invalid request body: {}", e.body_text())))
565}
566
567pub fn lenient_body<T: serde::de::DeserializeOwned>(
570 r: Result<Json<serde_json::Value>, JsonRejection>,
571 known: &[&[&str]],
572) -> Result<(T, Vec<String>), ApiError> {
573 use crate::fields;
574 let value = body(r)?;
575 let mut ignored = Vec::new();
576 if let Some(obj) = value.as_object() {
577 ignored = fields::unknown(obj, known);
578 if let Some(nested) = obj.get("scrapeOptions").and_then(|v| v.as_object()) {
579 let scrape = [
580 fields::names::<crate::types::ScrapeOptions>(),
581 fields::CLIENT_IDENTITY,
582 ];
583 ignored.extend(
584 fields::unknown(nested, &scrape)
585 .into_iter()
586 .map(|k| format!("scrapeOptions.{k}")),
587 );
588 }
589 }
590 let parsed = serde_json::from_value(value)
591 .map_err(|e| ApiError::invalid(format!("invalid request body: {e}")))?;
592 Ok((parsed, ignored))
593}
594
595pub fn base_url(s: &AppState, headers: &HeaderMap) -> String {
596 if let Some(u) = &s.public_url {
597 return u.clone();
598 }
599 let host = headers
600 .get(header::HOST)
601 .and_then(|v| v.to_str().ok())
602 .unwrap_or("localhost:3002");
603 format!("http://{host}")
604}
605
606pub fn webhook(
607 s: &AppState,
608 v: Option<&serde_json::Value>,
609) -> Result<Option<WebhookConfig>, ApiError> {
610 v.filter(|v| !v.is_null())
611 .map(|v| WebhookConfig::parse(v, s.webhooks.ssrf()))
612 .transpose()
613}
614
615async fn render_metrics(State(s): State<AppState>) -> impl IntoResponse {
616 let mut out = s.metrics.render(&s.jobs.counts().await);
617 if let Some(extra) = &s.extra_metrics {
618 out.push_str(&extra());
619 }
620 (
621 StatusCode::OK,
622 [(header::CONTENT_TYPE, "text/plain; version=0.0.4")],
623 out,
624 )
625}
626
627#[cfg(test)]
628mod batch_concurrency_tests {
629 use super::batch_concurrency;
630
631 #[test]
634 fn admission_decides_and_the_user_caps() {
635 assert_eq!(batch_concurrency(None, None, None), 8);
636 assert_eq!(batch_concurrency(Some(20), None, None), 8);
637 assert_eq!(batch_concurrency(Some(3), None, None), 3);
638 assert_eq!(batch_concurrency(None, None, Some(32)), 32);
639 assert_eq!(batch_concurrency(Some(64), None, Some(32)), 32);
640 assert_eq!(batch_concurrency(None, Some(50), None), 50);
641 assert_eq!(batch_concurrency(None, Some(50), Some(32)), 32);
642 assert_eq!(batch_concurrency(Some(5), Some(50), None), 5);
643 assert_eq!(batch_concurrency(Some(0), Some(0), None), 1);
644 }
645}