pagesift-api 0.90.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，
/// 每段一产出就写，不拼成整块。客户端中途断开就停止抓取。响应头发出后才失败的，用 `error` 事件报告。
fn stream_scrape(
    s: AppState,
    caller: Option<Caller>,
    url: String,
    opts: ScrapeOptions,
) -> Response {
    let (tx, body) = crate::stream_body::channel();
    tokio::spawn(async move {
        let r = tokio::select! {
            r = s.engine.scrape(&url, &opts) => r,
            () = tx.closed() => 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)
}