收口 MinerU OCR 任务链与上下文

2026-06-01 10:30:42 +08:00
parent 1481de2018
commit 610b23d3a5
13 changed files with 710 additions and 55 deletions
@@ -18,7 +18,7 @@ use std::convert::Infallible;
 use std::path::PathBuf;
 use std::pin::Pin;
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
-use tokio::sync::broadcast::error::RecvError;
+use tokio::sync::broadcast::{error::RecvError, Receiver};
 use tokio::time::timeout;

 type BoxedEventStream =
@@ -80,6 +80,7 @@ async fn build_document_events_stream(
        .local_folder_watcher_registry()
        .subscribe(&canonical_root)
        .map_err(|error| WebError::internal(error).with_context(&context))?;
+    let local_ocr_job_rx = state.local_ocr_job_tx.subscribe();

    let initial = json!({
        "sourceKind": "local_folder",
@@ -89,29 +90,79 @@ async fn build_document_events_stream(
        "revision": system_time_ms(SystemTime::now()),
    });
    let stream = stream::unfold(
-        (Some(initial), subscription, document_relative_path),
-        |(initial, mut subscription, document_relative_path)| async move {
+        (
+            Some(initial),
+            subscription,
+            document_relative_path,
+            query.root_uri.clone(),
+            local_ocr_job_rx,
+        ),
+        |(
+            initial,
+            mut subscription,
+            document_relative_path,
+            root_uri,
+            mut local_ocr_job_rx,
+        )| async move {
            if let Some(payload) = initial {
                return Some((
                    Ok(stream_event("ready", &payload)),
-                    (None, subscription, document_relative_path),
+                    (
+                        None,
+                        subscription,
+                        document_relative_path,
+                        root_uri,
+                        local_ocr_job_rx,
+                    ),
                ));
            }
            loop {
-                match subscription.receiver.recv().await {
-                    Ok(payload) => {
-                        if let Some(expected) = document_relative_path.as_deref() {
-                            if !document_event_targets_relative_path(&payload, expected) {
-                                continue;
+                tokio::select! {
+                    watcher_result = subscription.receiver.recv() => {
+                        match watcher_result {
+                            Ok(payload) => {
+                                if let Some(expected) = document_relative_path.as_deref() {
+                                    if !document_event_targets_relative_path(&payload, expected) {
+                                        continue;
+                                    }
+                                }
+                                return Some((
+                                    Ok(stream_event("change", &payload)),
+                                    (
+                                        None,
+                                        subscription,
+                                        document_relative_path,
+                                        root_uri,
+                                        local_ocr_job_rx,
+                                    ),
+                                ));
                            }
+                            Err(RecvError::Lagged(_)) => continue,
+                            Err(RecvError::Closed) => return None,
+                        }
+                    }
+                    ocr_result = recv_matching_ocr_event(&mut local_ocr_job_rx, &root_uri) => {
+                        match ocr_result {
+                            Some(payload) => {
+                                if let Some(expected) = document_relative_path.as_deref() {
+                                    if !document_event_targets_relative_path(&payload, expected) {
+                                        continue;
+                                    }
+                                }
+                                return Some((
+                                    Ok(stream_event("local_ocr.job.updated", &payload)),
+                                    (
+                                        None,
+                                        subscription,
+                                        document_relative_path,
+                                        root_uri,
+                                        local_ocr_job_rx,
+                                    ),
+                                ));
+                            }
+                            None => continue,
                        }
-                        return Some((
-                            Ok(stream_event("change", &payload)),
-                            (None, subscription, document_relative_path),
-                        ));
                    }
-                    Err(RecvError::Lagged(_)) => continue,
-                    Err(RecvError::Closed) => return None,
                }
            }
        },
@@ -121,6 +172,20 @@ async fn build_document_events_stream(
    Ok((HeaderMap::new(), stream))
 }

+async fn recv_matching_ocr_event(rx: &mut Receiver<Value>, root_uri: &str) -> Option<Value> {
+    loop {
+        match rx.recv().await {
+            Ok(payload) => {
+                if payload.get("rootUri").and_then(Value::as_str) == Some(root_uri) {
+                    return Some(payload);
+                }
+            }
+            Err(RecvError::Lagged(_)) => continue,
+            Err(RecvError::Closed) => return None,
+        }
+    }
+}
+
 /// Build the tree live stream: emits `snapshot` (initial) and `resync` (on watcher change)
 /// with full sidebar + file tree projections.
 ///
@@ -181,9 +181,11 @@ pub(crate) async fn create_job(
        body.force,
    )?;
    let now = now_ms();
+    let mut entry = build_index_entry(&plan, "queued", now, "", None);
+    upsert_and_broadcast_ocr_index_entry(&state, &root, root_uri, entry.clone())?;
    if let Some(error) = body.mock_error.as_deref() {
-        let entry = build_index_entry(&plan, "failed", now, "", Some(redact_error(error)));
-        upsert_ocr_index_entry(&root, entry.clone())?;
+        entry = advance_ocr_entry(&entry, "failed", now_ms(), "", Some(redact_error(error)));
+        upsert_and_broadcast_ocr_index_entry(&state, &root, root_uri, entry.clone())?;
        return Ok(ok_json(
            &context,
            json!({
@@ -193,27 +195,44 @@ pub(crate) async fn create_job(
        ));
    }
    let markdown = if provider == "mock" {
+        entry = advance_ocr_entry(&entry, "writing_sidecar", now_ms(), "", None);
+        upsert_and_broadcast_ocr_index_entry(&state, &root, root_uri, entry.clone())?;
        body.mock_markdown
            .unwrap_or_else(|| format!("OCR mock result for {}", plan.source_root_relative_path))
    } else {
-        match run_mineru_ocr(&root, &plan, token.unwrap_or_default()).await {
+        entry = advance_ocr_entry(&entry, "uploading", now_ms(), "", None);
+        upsert_and_broadcast_ocr_index_entry(&state, &root, root_uri, entry.clone())?;
+        match run_mineru_ocr(
+            &state,
+            &root,
+            root_uri,
+            &plan,
+            token.unwrap_or_default(),
+            entry.clone(),
+        )
+        .await
+        {
            Ok(markdown) => markdown,
            Err(error) => {
-                let failed_entry = build_index_entry(
-                    &plan,
+                let failed_entry = advance_ocr_entry(
+                    &entry,
                    "failed",
-                    now,
+                    now_ms(),
                    "",
                    Some(redact_error(error.message())),
                );
-                upsert_ocr_index_entry(&root, failed_entry)?;
+                upsert_and_broadcast_ocr_index_entry(&state, &root, root_uri, failed_entry)?;
                return Err(error.with_context(&context));
            }
        }
    };
+    if entry.status != "writing_sidecar" {
+        entry = advance_ocr_entry(&entry, "writing_sidecar", now_ms(), "", None);
+        upsert_and_broadcast_ocr_index_entry(&state, &root, root_uri, entry.clone())?;
+    }
    write_ocr_sidecar(&plan, &markdown, now)?;
-    let entry = build_index_entry(&plan, "done", now, &markdown, None);
-    upsert_ocr_index_entry(&root, entry.clone())?;
+    entry = advance_ocr_entry(&entry, "done", now_ms(), &markdown, None);
+    upsert_and_broadcast_ocr_index_entry(&state, &root, root_uri, entry.clone())?;
    Ok(ok_json(
        &context,
        json!({
@@ -463,9 +482,12 @@ pub(crate) fn strip_ocr_frontmatter(markdown: &str) -> &str {
 }

 async fn run_mineru_ocr(
+    state: &AppState,
    root: &Path,
+    root_uri: &str,
    plan: &OcrSidecarPlan,
    token: String,
+    mut entry: OcrIndexEntry,
 ) -> Result<String, WebError> {
    let config = MineruClientConfig {
        api_base_url: mineru_api_base_url(),
@@ -491,7 +513,11 @@ async fn run_mineru_ocr(
    let client = reqwest::Client::new();
    let upload = create_mineru_upload_task(&client, &config, file_name).await?;
    upload_mineru_source(&client, &upload.upload_url, bytes).await?;
+    entry = advance_ocr_entry(&entry, "mineru_processing", now_ms(), "", None);
+    upsert_and_broadcast_ocr_index_entry(state, root, root_uri, entry.clone())?;
    let zip_url = poll_mineru_result_zip_url(&client, &config, &upload.batch_id).await?;
+    entry = advance_ocr_entry(&entry, "downloading", now_ms(), "", None);
+    upsert_and_broadcast_ocr_index_entry(state, root, root_uri, entry)?;
    let zip_bytes = download_mineru_result_zip(&client, &zip_url).await?;
    extract_mineru_markdown_from_zip(&zip_bytes)
 }
@@ -953,6 +979,23 @@ fn build_index_entry(
    }
 }

+fn advance_ocr_entry(
+    previous: &OcrIndexEntry,
+    status: &str,
+    now: u128,
+    markdown: &str,
+    error: Option<String>,
+) -> OcrIndexEntry {
+    let mut next = previous.clone();
+    next.status = status.to_string();
+    next.updated_at_ms = now;
+    if !markdown.is_empty() {
+        next.plain_text_preview = markdown.chars().take(240).collect();
+    }
+    next.error = error;
+    next
+}
+
 fn read_ocr_index(root: &Path) -> Result<OcrIndex, WebError> {
    let path = ocr_index_path(root);
    if !path.exists() {
@@ -1011,6 +1054,42 @@ fn upsert_ocr_index_entry(root: &Path, entry: OcrIndexEntry) -> Result<(), WebEr
    write_ocr_index(root, &index)
 }

+fn upsert_and_broadcast_ocr_index_entry(
+    state: &AppState,
+    root: &Path,
+    root_uri: &str,
+    entry: OcrIndexEntry,
+) -> Result<(), WebError> {
+    upsert_ocr_index_entry(root, entry.clone())?;
+    broadcast_ocr_job_update(state, root, root_uri, &entry);
+    Ok(())
+}
+
+fn broadcast_ocr_job_update(state: &AppState, root: &Path, root_uri: &str, entry: &OcrIndexEntry) {
+    let job = ocr_job_payload(root, entry);
+    let key = format!("{root_uri}:{}", entry.source_root_relative_path);
+    if matches!(entry.status.as_str(), "done" | "failed" | "interrupted") {
+        if let Ok(mut jobs) = state.local_ocr_active_jobs.write() {
+            jobs.remove(&key);
+        }
+    } else if let Ok(mut jobs) = state.local_ocr_active_jobs.write() {
+        jobs.insert(key, job.clone());
+    }
+    let payload = json!({
+        "schema": "mnote.local_ocr.job.updated.v1",
+        "kind": "local_ocr_job_updated",
+        "eventType": "local_ocr.job.updated",
+        "sourceKind": "local_folder",
+        "rootUri": root_uri,
+        "relativePath": entry.source_root_relative_path,
+        "documentId": entry.owner_document_id,
+        "revision": entry.updated_at_ms.to_string(),
+        "job": job,
+    });
+    let _ = state.local_ocr_job_tx.send(payload.clone());
+    let _ = state.stream_delta_tx.send(payload);
+}
+
 fn ocr_job_payload(root: &Path, entry: &OcrIndexEntry) -> Value {
    let stale = source_is_stale(root, entry);
    json!({
@@ -1267,8 +1346,8 @@ mod tests {
        root
    }

-    fn app() -> axum::Router {
-        build_app(AppState::new(AppConfig {
+    fn test_app_state() -> AppState {
+        AppState::new(AppConfig {
            service_name: "mnote-web".into(),
            service_version: "0.1.0".into(),
            bind_addr: "127.0.0.1:0".into(),
@@ -1287,7 +1366,15 @@ mod tests {
            dev_user_id: "dev-user".into(),
            dev_user_name: "开发用户".into(),
            dev_user_email: "dev@mnote.local".into(),
-        }))
+        })
+    }
+
+    fn app() -> axum::Router {
+        build_app(test_app_state())
+    }
+
+    fn app_with_state(state: AppState) -> axum::Router {
+        build_app(state)
    }

    fn write_workspace_manifest(root: &Path) {
@@ -1501,6 +1588,70 @@ mod tests {
        let _ = fs::remove_dir_all(root);
    }

+    #[tokio::test]
+    async fn local_ocr_jobs_route_broadcasts_stage_updates() {
+        let root = temp_root("mnote-local-ocr-events");
+        write_workspace_manifest(&root);
+        fs::write(root.join("docs").join("Page.md"), "# Page\n").expect("page");
+        fs::write(
+            root.join("docs").join("Page.assets").join("photo.png"),
+            b"png",
+        )
+        .expect("photo");
+        let root_uri = format!("file://{}", root.display());
+        let state = test_app_state();
+        let mut events = state.local_ocr_job_tx.subscribe();
+        let response = app_with_state(state.clone())
+            .oneshot(
+                Request::builder()
+                    .method("POST")
+                    .uri("/api/local-folder/ocr/jobs")
+                    .header("content-type", "application/json")
+                    .header("x-mnote-actor-id", "user_test")
+                    .header("x-mnote-actor-type", "user")
+                    .body(Body::from(
+                        json!({
+                            "rootUri": root_uri,
+                            "documentId": "local-md:docs~2FPage.md",
+                            "sourceRootRelativePath": "docs/Page.assets/photo.png",
+                            "provider": "mock",
+                            "mockMarkdown": "Event OCR Token"
+                        })
+                        .to_string(),
+                    ))
+                    .expect("request"),
+            )
+            .await
+            .expect("response");
+        assert_eq!(response.status(), StatusCode::OK);
+
+        let mut statuses = Vec::new();
+        for _ in 0..3 {
+            let payload = tokio::time::timeout(Duration::from_secs(1), events.recv())
+                .await
+                .expect("ocr event timeout")
+                .expect("ocr event");
+            assert_eq!(payload["eventType"].as_str(), Some("local_ocr.job.updated"));
+            assert_eq!(
+                payload["job"]["sourceRootRelativePath"].as_str(),
+                Some("docs/Page.assets/photo.png")
+            );
+            statuses.push(
+                payload["job"]["status"]
+                    .as_str()
+                    .unwrap_or_default()
+                    .to_string(),
+            );
+        }
+        assert_eq!(statuses, vec!["queued", "writing_sidecar", "done"]);
+        assert!(state
+            .local_ocr_active_jobs
+            .read()
+            .expect("active jobs lock")
+            .is_empty());
+        let _ = fs::remove_dir_all(root);
+    }
+
    #[tokio::test]
    async fn local_ocr_jobs_route_rejects_missing_mineru_token() {
        let old_mnote_token = std::env::var("MNOTE_MINERU_API_TOKEN").ok();