1use 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
25pub const PAGE_SIZE: usize = 100;
27pub const DEFAULT_BATCH_CONCURRENCY: usize = 8;
28
29#[derive(Clone, Default)]
30pub struct Config {
31 pub api_keys: Vec<String>,
33 pub public_url: Option<String>,
35 pub max_batch_concurrency: Option<usize>,
36 pub store: Option<Store>,
38 pub key_limits: HashMap<String, KeyLimits>,
40 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
80pub 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
153async 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
189async 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}