pagesift-api 0.91.0

//! Jina Reader protocol handler (`GET /r/{*url}` and `POST /r`).
//! Provides direct Markdown extraction and JSON responses for AI Agents.

use axum::Json;
use axum::extract::{Extension, Path, State};
use axum::http::{HeaderMap, StatusCode, header};
use axum::response::{IntoResponse, Response};
use serde::Deserialize;
use std::sync::atomic::Ordering;

use super::stream_sse::{error_event, sse_events, sse_response};
use crate::limits::Caller;
use crate::metrics::PageSource;
use crate::server::AppState;
use crate::types::ScrapeOptions;

#[derive(Debug, Deserialize)]
pub struct JinaPostRequest {
    pub url: String,
}

/// Handler for `GET /r/{*target_url}`.
pub(crate) async fn handle_jina_get(
    State(s): State<AppState>,
    caller: Option<Extension<Caller>>,
    headers: HeaderMap,
    Path(raw_url): Path<String>,
) -> Response {
    let clean_url = parse_target_url(&raw_url);
    process_jina_scrape(s, caller, headers, &clean_url).await
}

/// Handler for `POST /r`.
pub(crate) async fn handle_jina_post(
    State(s): State<AppState>,
    caller: Option<Extension<Caller>>,
    headers: HeaderMap,
    Json(payload): Json<JinaPostRequest>,
) -> Response {
    let clean_url = parse_target_url(&payload.url);
    process_jina_scrape(s, caller, headers, &clean_url).await
}

fn parse_target_url(raw: &str) -> String {
    let trimmed = raw.trim();
    if trimmed.starts_with("http://") || trimmed.starts_with("https://") {
        trimmed.to_string()
    } else if let Some(idx) = trimmed.find("http:/") {
        if trimmed[idx..].starts_with("http://") {
            trimmed[idx..].to_string()
        } else {
            format!("http://{}", trimmed[idx + 6..].trim_start_matches('/'))
        }
    } else if let Some(idx) = trimmed.find("https:/") {
        if trimmed[idx..].starts_with("https://") {
            trimmed[idx..].to_string()
        } else {
            format!("https://{}", trimmed[idx + 7..].trim_start_matches('/'))
        }
    } else {
        format!("https://{}", trimmed.trim_start_matches('/'))
    }
}

async fn process_jina_scrape(
    s: AppState,
    caller: Option<Extension<Caller>>,
    headers: HeaderMap,
    target_url: &str,
) -> Response {
    if let Err(e) = crate::validate::check_url(target_url) {
        return e.into_response();
    }
    if let Some(Extension(c)) = &caller
        && let Err(e) = s.limiter.check_quota(c)
    {
        return e.into_response();
    }

    let opts = ScrapeOptions {
        only_main_content: true,
        ..Default::default()
    };
    let accept = headers
        .get(header::ACCEPT)
        .and_then(|v| v.to_str().ok())
        .unwrap_or("");
    if accept.contains("text/event-stream") {
        return stream_scrape(
            s,
            caller.map(|Extension(c)| c),
            target_url.to_string(),
            opts,
        );
    }

    match s.engine.scrape(target_url, &opts).await {
        Ok(doc) => {
            record_ok(&s, caller.as_ref().map(|Extension(c)| c)).await;
            let markdown = doc.markdown.unwrap_or_default();
            let title = doc.metadata.title.as_deref();

            if accept.contains("application/json") {
                let json_data = serde_json::json!({
                    "code": 200,
                    "status": 20000,
                    "data": {
                        "title": title.unwrap_or(""),
                        "description": doc.metadata.description.unwrap_or_default(),
                        "url": target_url,
                        "content": markdown,
                    }
                });
                Json(json_data).into_response()
            } else {
                // Default: Direct Raw Markdown for LLM Agents
                Response::builder()
                    .header(header::CONTENT_TYPE, "text/markdown; charset=utf-8")
                    .body(axum::body::Body::from(markdown))
                    .unwrap_or_else(|_| {
                        (StatusCode::INTERNAL_SERVER_ERROR, "response builder error")
                            .into_response()
                    })
            }
        }
        Err(e) => {
            s.metrics.scrape_err.fetch_add(1, Ordering::Relaxed);
            e.into_response()
        }
    }
}

async fn record_ok(s: &AppState, caller: Option<&Caller>) {
    s.metrics.scrape_ok.fetch_add(1, Ordering::Relaxed);
    s.metrics.page(PageSource::Scrape);
    if let Some(c) = caller {
        s.limiter.consume_pages(c, 1).await;
    }
}

/// SSE：响应头立刻发出，抓取在流里进行；整页抓完、提取完之后才逐个写出 progress、各段 content 和 done
/// （按段分帧写出，不拼成整块），抓取期间按心跳间隔写注释保活。不边提取边推：正文识别要看完整个 DOM
/// 才知道哪些块是正文（打分、去模板），提取本身只占请求耗时的一小部分（0.91 发布说明里有本机实测），
/// 第一个 `content` 的时刻等于整页提取完成的时刻（0.90 的发布说明写成了"一产出就写"，不准确）。
/// 客户端中途断开就停止抓取。响应头发出后才失败的，用 `error` 事件报告。
fn stream_scrape(
    s: AppState,
    caller: Option<Caller>,
    url: String,
    opts: ScrapeOptions,
) -> Response {
    stream_scrape_with(s, caller, url, opts, crate::firecrawl::stream::KEEPALIVE)
}

/// 同 [`stream_scrape`]，`keepalive` 是抓取期间写心跳注释（`: keepalive`）的间隔：渲染、慢站点
/// 可能几十秒才有结果，期间不写字节，代理和客户端会按空闲超时断开（0.90 如此）。SSE 客户端忽略注释行，
/// 事件本身的字节不变。
fn stream_scrape_with(
    s: AppState,
    caller: Option<Caller>,
    url: String,
    opts: ScrapeOptions,
    keepalive: std::time::Duration,
) -> Response {
    let (tx, body) = crate::stream_body::channel();
    tokio::spawn(async move {
        let scrape = s.engine.scrape(&url, &opts);
        tokio::pin!(scrape);
        let mut tick = tokio::time::interval_at(tokio::time::Instant::now() + keepalive, keepalive);
        let r = loop {
            tokio::select! {
                r = &mut scrape => break r,
                () = tx.closed() => return,
                _ = tick.tick() => {
                    if tx.send(axum::body::Bytes::from_static(b": keepalive\n\n")).await.is_err() {
                        return;
                    }
                }
            }
        };
        match r {
            Ok(doc) => {
                record_ok(&s, caller.as_ref()).await;
                let markdown = doc.markdown.unwrap_or_default();
                for event in sse_events(&url, doc.metadata.title.as_deref(), &markdown) {
                    if tx.send(event.into()).await.is_err() {
                        return;
                    }
                }
            }
            Err(e) => {
                s.metrics.scrape_err.fetch_add(1, Ordering::Relaxed);
                let _ = tx.send(error_event(&e).into()).await;
            }
        }
    });
    sse_response(body)
}