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

use crate::app::AppState;
use crate::context::RequestContext;
use crate::error::WebError;
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::http::{HeaderMap, HeaderName, HeaderValue};
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 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?;

            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"),
    ))
}

pub async fn tree_events(
    State(state): State<AppState>,
    Extension(context): Extension<RequestContext>,
    Query(query): Query<StreamSnapshotQuery>,
) -> Result<
    (
        HeaderMap,
        Sse<impl futures_util::Stream<Item = Result<Event, Infallible>>>,
    ),
    WebError,
> {
    let mut headers = HeaderMap::new();
    if let Ok(name) = HeaderName::from_lowercase(b"x-mnote-web-owner") {
        headers.insert(name, HeaderValue::from_static("mnote-web"));
    }
    if let Ok(name) = HeaderName::from_lowercase(b"x-mnote-tree-stream-owner") {
        headers.insert(name, HeaderValue::from_static("rust-web"));
    }
    let sse = events(State(state), Extension(context), Query(query)).await?;
    Ok((headers, sse))
}

#[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 {
    let event_id = payload
        .get("revision")
        .and_then(|value| {
            value
                .as_str()
                .map(ToOwned::to_owned)
                .or_else(|| value.as_u64().map(|number| number.to_string()))
        })
        .or_else(|| {
            payload
                .get("cursor")
                .and_then(Value::as_str)
                .map(ToOwned::to_owned)
        })
        .unwrap_or_else(|| "0".into());
    Event::default()
        .event(event_name)
        .id(event_id)
        .json_data(payload)
        .expect("SSE 事件必须可序列化")
}

#[cfg(test)]
mod tests {
    use crate::app::{build_app, AppConfig, AppState};
    use axum::body::{to_bytes, Body};
    use axum::http::{Request, StatusCode};
    use tower::util::ServiceExt;

    fn app() -> axum::Router {
        build_app(AppState::new(AppConfig {
            service_name: "mnote-web".into(),
            service_version: "0.1.0".into(),
            bind_addr: "127.0.0.1:0".into(),
            public_bind_addr: "127.0.0.1:3000".into(),
            legacy_next_base_url: Some("http://127.0.0.1:3100".into()),
            enable_legacy_next_compat: true,
            enable_debug_shell_routes: false,
            hermes_base_path: "/api/hermes".into(),
            compat_next_base_path: "/api/compat/next".into(),
            convex_url: None,
            convex_admin_key: None,
            allow_dev_fixtures: true,
            query_fixtures_json: Some(r#"{"sidebar:datasetList":{"active_workspace_id":"ws_demo","documents":[{"id":"page_root","workspace_id":"ws_demo","title":"工作区首页","parent_id":null,"sort_order":0,"is_starred":true,"is_template":false,"created_at":"2026-04-16T00:00:00Z","updated_at":"2026-04-16T00:00:00Z"},{"id":"page_child","workspace_id":"ws_demo","title":"子页面","parent_id":"page_root","sort_order":1,"is_starred":false,"is_template":false,"created_at":"2026-04-16T00:00:00Z","updated_at":"2026-04-16T00:00:00Z"}],"trashed_documents":[],"media_assets":[],"trashed_media_assets":[],"mindmap_assets":[],"trashed_mindmap_assets":[],"table_assets":[],"trashed_table_assets":[],"mindmap_docs":[],"mindmap_asset_children":{}},"bridgeLogs:listWorkspaceOverview":{"workspace_id":"ws_demo","command_logs":[{"command_id":"cmd_1","request_id":"req_1","status":"applied","created_at":"2026-04-16T00:00:00Z"}],"domain_events":[{"command_id":"cmd_1","status":"published","created_at":"2026-04-16T00:00:00Z"}],"next_cursor":"cursor_demo","has_more":false,"filters":{"command_status":null,"event_status":null,"target_page_id":null,"target_block_id":null,"aggregate_type":null,"aggregate_id":null},"generated_at":"2026-04-16T00:00:00Z"}}"#.into()),
            mutation_fixtures_json: None,
            dev_user_id: "dev-user".into(),
            dev_user_name: "开发用户".into(),
            dev_user_email: "dev@mnote.local".into(),
        }))
    }

    #[tokio::test]
    async fn sse_route_returns_workspace_snapshot_event() {
        let response = app()
            .oneshot(
                Request::builder()
                    .uri("/api/stream/events?workspaceId=ws_demo&maxPolls=0")
                    .body(Body::empty())
                    .expect("request"),
            )
            .await
            .expect("response");

        assert_eq!(response.status(), StatusCode::OK);
        let body = to_bytes(response.into_body(), usize::MAX)
            .await
            .expect("body");
        let text = String::from_utf8(body.to_vec()).expect("utf8");
        assert!(text.contains("event: snapshot") || text.contains("event:snapshot"));
        assert!(text.contains("\"stream\":\"workspace\""));
        assert!(text.contains("\"projection\":\"sidebar_tree\""));
        assert!(text.contains("\"workspaceId\":\"ws_demo\""));
    }

    #[tokio::test]
    async fn tree_realtime_route_returns_rust_web_owned_snapshot_event() {
        let response = app()
            .oneshot(
                Request::builder()
                    .uri("/api/tree/events?workspaceId=ws_demo&maxPolls=0")
                    .body(Body::empty())
                    .expect("request"),
            )
            .await
            .expect("response");

        assert_eq!(response.status(), StatusCode::OK);
        assert_eq!(
            response
                .headers()
                .get("x-mnote-web-owner")
                .and_then(|value| value.to_str().ok()),
            Some("mnote-web")
        );
        assert_eq!(
            response
                .headers()
                .get("x-mnote-tree-stream-owner")
                .and_then(|value| value.to_str().ok()),
            Some("rust-web")
        );
        let body = to_bytes(response.into_body(), usize::MAX)
            .await
            .expect("body");
        let text = String::from_utf8(body.to_vec()).expect("utf8");
        assert!(text.contains("event: snapshot") || text.contains("event:snapshot"));
        assert!(text.contains("\"kind\":\"snapshot\""));
    }

    #[tokio::test]
    async fn tree_events_route_includes_event_id_and_revision() {
        let response = app()
            .oneshot(
                Request::builder()
                    .uri("/api/tree/events?workspaceId=ws_demo&maxPolls=0")
                    .body(Body::empty())
                    .expect("request"),
            )
            .await
            .expect("response");

        assert_eq!(response.status(), StatusCode::OK);
        let body = to_bytes(response.into_body(), usize::MAX)
            .await
            .expect("body");
        let text = String::from_utf8(body.to_vec()).expect("utf8");
        assert!(text.contains("id: "));
        assert!(text.contains("\"revision\""));
    }
}