Skip to main content

pagesift_api/
server.rs

1//! HTTP server routing: Firecrawl v1/v2 compatibility layers, Jina Reader adapter,
2//! health checks, and Prometheus metrics.
3
4use std::collections::HashMap;
5use std::sync::Arc;
6use std::time::Instant;
7
8use axum::extract::rejection::JsonRejection;
9use axum::extract::{Request, State};
10use axum::http::{HeaderMap, Method, StatusCode, header};
11use axum::middleware::{self, Next};
12use axum::response::{IntoResponse, Response};
13use axum::routing::get;
14use axum::{Json, Router};
15use pagesift_store::Store;
16use tokio::sync::Mutex;
17
18use crate::engine::Engine;
19use crate::error::{ApiError, ErrorCode};
20use crate::jobs::Jobs;
21use crate::limits::{Caller, KeyLimits, Limiter};
22use crate::metrics::Metrics;
23use crate::webhook::{WebhookConfig, WebhookSettings, Webhooks};
24
25/// Maximum documents per page when querying job status.
26pub const PAGE_SIZE: usize = 100;
27pub const DEFAULT_BATCH_CONCURRENCY: usize = 8;
28
29#[derive(Clone, Default)]
30pub struct Config {
31    /// When empty, authentication is disabled; otherwise requires `Authorization: Bearer <key>`.
32    pub api_keys: Vec<String>,
33    /// Public URL for building absolute resource links.
34    pub public_url: Option<String>,
35    pub max_batch_concurrency: Option<usize>,
36    /// Persistent store for jobs and resume capability.
37    pub store: Option<Store>,
38    /// Per-key custom limits.
39    pub key_limits: HashMap<String, KeyLimits>,
40    /// Default limits applied to other keys or unauthenticated callers.
41    pub default_limits: KeyLimits,
42    pub webhook: WebhookSettings,
43}
44
45impl std::fmt::Debug for Config {
46    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
47        f.debug_struct("Config")
48            .field("api_keys", &self.api_keys.len())
49            .field("public_url", &self.public_url)
50            .field("max_batch_concurrency", &self.max_batch_concurrency)
51            .field("store", &self.store.is_some())
52            .field("key_limits", &self.key_limits.len())
53            .field("default_limits", &self.default_limits)
54            .field("webhook", &self.webhook)
55            .finish()
56    }
57}
58
59#[derive(Clone)]
60pub struct AppState {
61    pub engine: Arc<dyn Engine>,
62    pub keys: Arc<HashMap<String, Arc<KeyLimits>>>,
63    pub default_limits: Arc<KeyLimits>,
64    pub public_url: Option<String>,
65    pub batch_concurrency: usize,
66    pub jobs: Jobs,
67    pub limiter: Arc<Limiter>,
68    pub webhooks: Webhooks,
69    pub extracts: Arc<Mutex<HashMap<String, ExtractJob>>>,
70    pub metrics: Arc<Metrics>,
71}
72
73pub struct ExtractJob {
74    pub status: &'static str,
75    pub data: Option<serde_json::Value>,
76    pub error: Option<ApiError>,
77    pub created: std::time::SystemTime,
78}
79
80/// Construct complete HTTP router with Firecrawl v1/v2, Jina Reader, and metrics.
81pub fn router(engine: Arc<dyn Engine>, config: Config) -> Router {
82    let default_limits = Arc::new(config.default_limits);
83    let mut keys: HashMap<String, Arc<KeyLimits>> = config
84        .api_keys
85        .into_iter()
86        .map(|k| (k, default_limits.clone()))
87        .collect();
88    keys.extend(config.key_limits.into_iter().map(|(k, l)| (k, Arc::new(l))));
89
90    let accounts: HashMap<String, Arc<KeyLimits>> = keys
91        .iter()
92        .map(|(k, l)| {
93            (
94                Caller::new(Some(k), l.clone()).account.to_string(),
95                l.clone(),
96            )
97        })
98        .collect();
99
100    let metrics = Arc::new(Metrics::default());
101    let limiter = Arc::new(Limiter::new(config.store.clone()));
102    let webhooks = Webhooks::new(config.webhook, metrics.clone());
103    let jobs = Jobs::new(
104        config.store,
105        limiter.clone(),
106        Some(webhooks.clone()),
107        metrics.clone(),
108    );
109
110    if tokio::runtime::Handle::try_current().is_ok() {
111        let fallback = default_limits.clone();
112        jobs.restore(&engine, &|account| {
113            accounts
114                .get(account)
115                .cloned()
116                .unwrap_or_else(|| fallback.clone())
117        });
118    } else {
119        tracing::warn!("router built outside a tokio runtime: saved jobs are not resumed");
120    }
121
122    let state = AppState {
123        engine: engine.clone(),
124        keys: Arc::new(keys),
125        default_limits,
126        public_url: config
127            .public_url
128            .map(|u| u.trim_end_matches('/').to_string()),
129        batch_concurrency: config
130            .max_batch_concurrency
131            .unwrap_or(DEFAULT_BATCH_CONCURRENCY),
132        jobs,
133        limiter,
134        webhooks,
135        extracts: Arc::default(),
136        metrics,
137    };
138
139    let firecrawl = crate::firecrawl::firecrawl_routes()
140        .route_layer(middleware::from_fn_with_state(state.clone(), auth));
141    let jina = crate::jina::jina_router(engine);
142
143    Router::new()
144        .route("/healthz", get(|| async { "ok" }))
145        .route("/metrics", get(render_metrics))
146        .merge(firecrawl)
147        .with_state(state.clone())
148        .merge(jina)
149        .fallback(|| async { ApiError::new(ErrorCode::NotFound, "no such endpoint") })
150        .layer(middleware::from_fn_with_state(state.clone(), observe))
151}
152
153/// Authentication and rate limiting middleware.
154async fn auth(
155    State(s): State<AppState>,
156    headers: HeaderMap,
157    mut req: Request,
158    next: Next,
159) -> Response {
160    let caller = if s.keys.is_empty() {
161        Caller::new(None, s.default_limits.clone())
162    } else {
163        let key = headers
164            .get(header::AUTHORIZATION)
165            .and_then(|v| v.to_str().ok())
166            .and_then(|v| v.strip_prefix("Bearer "))
167            .map(str::trim);
168        match key.and_then(|k| s.keys.get(k).map(|l| (k, l))) {
169            Some((k, l)) => Caller::new(Some(k), l.clone()),
170            None => {
171                return ApiError::new(ErrorCode::Unauthorized, "missing or invalid API key")
172                    .into_response();
173            }
174        }
175    };
176    if req.method() == Method::POST
177        && let Err(e) = s.limiter.check_rate(&caller)
178    {
179        let mut resp = e.into_response();
180        resp.extensions_mut().insert(caller);
181        return resp;
182    }
183    req.extensions_mut().insert(caller.clone());
184    let mut resp = next.run(req).await;
185    resp.extensions_mut().insert(caller);
186    resp
187}
188
189/// Metrics and structured request tracing middleware.
190async fn observe(State(s): State<AppState>, req: Request, next: Next) -> Response {
191    let route = req
192        .extensions()
193        .get::<axum::extract::MatchedPath>()
194        .map(|p| p.as_str().to_string())
195        .unwrap_or_else(|| "unmatched".into());
196    let method = req.method().clone();
197    let started = Instant::now();
198    s.metrics.request_started();
199    let resp = next.run(req).await;
200    let elapsed = started.elapsed();
201    let status = resp.status().as_u16();
202    s.metrics.request_finished(&route, status, elapsed);
203    let code = resp.extensions().get::<ErrorCode>().map(|c| c.as_str());
204    if let Some(code) = code {
205        s.metrics.error(code, "http");
206    }
207    let account = resp.extensions().get::<Caller>().map(|c| c.account.clone());
208    let latency_ms = elapsed.as_secs_f64() * 1000.0;
209    if route.starts_with("/v2/") || route.starts_with("/v1/") || route.starts_with("/r") {
210        tracing::info!(
211            target: "pagesift::http",
212            %method, %route, status, latency_ms, account = account.as_deref(), code,
213            "request"
214        );
215    } else {
216        tracing::debug!(target: "pagesift::http", %method, %route, status, latency_ms, "request");
217    }
218    resp
219}
220
221pub fn body<T>(r: Result<Json<T>, JsonRejection>) -> Result<T, ApiError> {
222    r.map(|Json(v)| v)
223        .map_err(|e| ApiError::invalid(format!("invalid request body: {}", e.body_text())))
224}
225
226pub fn base_url(s: &AppState, headers: &HeaderMap) -> String {
227    if let Some(u) = &s.public_url {
228        return u.clone();
229    }
230    let host = headers
231        .get(header::HOST)
232        .and_then(|v| v.to_str().ok())
233        .unwrap_or("localhost:3002");
234    format!("http://{host}")
235}
236
237pub fn webhook(
238    s: &AppState,
239    v: Option<&serde_json::Value>,
240) -> Result<Option<WebhookConfig>, ApiError> {
241    v.filter(|v| !v.is_null())
242        .map(|v| WebhookConfig::parse(v, s.webhooks.ssrf()))
243        .transpose()
244}
245
246async fn render_metrics(State(s): State<AppState>) -> impl IntoResponse {
247    let out = s.metrics.render(&s.jobs.counts().await);
248    (
249        StatusCode::OK,
250        [(header::CONTENT_TYPE, "text/plain; version=0.0.4")],
251        out,
252    )
253}