4-26 树rust-2

rust/crates/mnote-web/src/routes/sse.rs

@@ -1,34 +1,138 @@
 use crate::app::AppState;
 use crate::context::RequestContext;
 use crate::error::WebError;
-use crate::routes::stream_support::{load_stream_snapshot, StreamSnapshotQuery};
+use crate::routes::stream_support::{
+    build_stream_delta_payload, load_stream_overview, load_stream_snapshot,
+    read_stream_cursor_from_payload, resolve_stream_change, with_stream_kind, StreamChangeKind,
+    StreamSnapshotQuery,
+};
 use axum::extract::{Extension, Query, State};
 use axum::response::sse::{Event, KeepAlive, Sse};
 use futures_util::stream;
 use serde_json::Value;
 use std::convert::Infallible;
 use std::time::Duration;
+use tokio::time::sleep;
 
 pub async fn events(
     State(state): State<AppState>,
     Extension(context): Extension<RequestContext>,
     Query(query): Query<StreamSnapshotQuery>,
 ) -> Result<Sse<impl futures_util::Stream<Item = Result<Event, Infallible>>>, WebError> {
-    let payload = load_stream_snapshot(state.config(), &context, &query).await?;
-    let event = snapshot_event(&payload);
+    let initial_payload = load_stream_snapshot(state.config(), &context, &query).await?;
+    let initial_cursor = read_stream_cursor_from_payload(&initial_payload);
+    let max_polls = query.max_polls;
+    let poll_ms = query.poll_ms.unwrap_or(2_000).max(250);
+    let state_for_stream = state.clone();
+    let context_for_stream = context.clone();
+    let query_for_stream = query.clone();
+    let stream = stream::unfold(
+        Some(StreamPollState {
+            app_state: state_for_stream,
+            context: context_for_stream,
+            query: query_for_stream,
+            current_cursor: initial_cursor,
+            polls: 0,
+            initial_payload,
+            initial_emitted: false,
+        }),
+        move |state| async move {
+            let mut state = state?;
 
-    Ok(Sse::new(stream::iter(vec![Ok(event)])).keep_alive(
+            if !state.initial_emitted {
+                state.initial_emitted = true;
+                return Some((
+                    Ok(stream_event("snapshot", &state.initial_payload)),
+                    Some(state),
+                ));
+            }
+
+            loop {
+                if let Some(max_polls) = max_polls {
+                    if state.polls >= max_polls {
+                        return None;
+                    }
+                }
+                state.polls += 1;
+                sleep(Duration::from_millis(poll_ms)).await;
+
+                let Ok((workspace_id, overview)) =
+                    load_stream_overview(state.app_state.config(), &state.context, &state.query)
+                        .await
+                else {
+                    return None;
+                };
+                let Some(change) =
+                    resolve_stream_change(&overview, state.current_cursor.as_deref())
+                else {
+                    continue;
+                };
+
+                state.current_cursor = change.cursor.clone();
+
+                match change.kind {
+                    StreamChangeKind::Delta => {
+                        let payload = build_stream_delta_payload(
+                            &state.context,
+                            &state.query,
+                            &workspace_id,
+                            &overview,
+                            change.cursor,
+                            change.delta.unwrap_or_else(|| serde_json::json!({ "op": "noop" })),
+                        );
+                        return Some((Ok(stream_event("delta", &payload)), Some(state)));
+                    }
+                    StreamChangeKind::Resync => {
+                        let mut next_query = state.query.clone();
+                        next_query.cursor = change.cursor;
+                        let Ok(snapshot_payload) = load_stream_snapshot(
+                            state.app_state.config(),
+                            &state.context,
+                            &next_query,
+                        )
+                        .await
+                        else {
+                            return None;
+                        };
+                        state.query = next_query;
+                        state.current_cursor =
+                            read_stream_cursor_from_payload(&snapshot_payload);
+                        return Some((
+                            Ok(stream_event(
+                                "resync",
+                                &with_stream_kind(&snapshot_payload, "resync"),
+                            )),
+                            Some(state),
+                        ));
+                    }
+                }
+            }
+        },
+    );
+
+    Ok(Sse::new(stream).keep_alive(
         KeepAlive::new()
             .interval(Duration::from_secs(15))
             .text("keepalive"),
     ))
 }
 
-fn snapshot_event(payload: &Value) -> Event {
+#[derive(Clone)]
+struct StreamPollState {
+    app_state: AppState,
+    context: RequestContext,
+    query: StreamSnapshotQuery,
+    current_cursor: Option<String>,
+    polls: u32,
+    initial_payload: Value,
+    initial_emitted: bool,
+}
+
+fn stream_event(event_name: &str, payload: &Value) -> Event {
     Event::default()
-        .event("snapshot")
+        .event(event_name)
         .json_data(payload)
-        .expect("SSE snapshot 事件必须可序列化")
+        .expect("SSE 事件必须可序列化")
 }
 
 #[cfg(test)]
@@ -62,7 +166,7 @@ mod tests {
         let response = app()
             .oneshot(
                 Request::builder()
-                    .uri("/api/stream/events?workspaceId=ws_demo")
+                    .uri("/api/stream/events?workspaceId=ws_demo&maxPolls=0")
                     .body(Body::empty())
                     .expect("request"),
             )
@@ -75,7 +179,8 @@ mod tests {
             .expect("body");
         let text = String::from_utf8(body.to_vec()).expect("utf8");
         assert!(text.contains("event: snapshot") || text.contains("event:snapshot"));
-        assert!(text.contains("\"scope\":\"workspace\""));
+        assert!(text.contains("\"stream\":\"workspace\""));
+        assert!(text.contains("\"projection\":\"sidebar_tree\""));
         assert!(text.contains("\"workspaceId\":\"ws_demo\""));
     }
 }