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, build_stream_push_delta_hint, 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::{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> {
    events_with_block_delta(state, context, query, None).await
}

async fn events_with_block_delta(
    state: AppState,
    context: RequestContext,
    query: StreamSnapshotQuery,
    block_delta_rx: Option<tokio::sync::broadcast::Receiver<Value>>,
) -> Result<Sse<impl futures_util::Stream<Item = Result<Event, Infallible>>>, WebError> {
    events_with_stream_delta(state, context, query, block_delta_rx, None).await
}

/// Unified SSE stream: when `stream_delta_rx` is present, broadcast-driven push takes priority;
/// polling acts as safety net. When absent, pure polling mode.
async fn events_with_stream_delta(
    state: AppState,
    context: RequestContext,
    query: StreamSnapshotQuery,
    block_delta_rx: Option<tokio::sync::broadcast::Receiver<Value>>,
    stream_delta_rx: Option<tokio::sync::broadcast::Receiver<Value>>,
) -> 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;
    // When push-driven, use long poll interval as safety net; otherwise normal polling
    let poll_ms = if stream_delta_rx.is_some() {
        query.poll_ms.unwrap_or(60_000).max(1_000)
    } else {
        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 subscription_workspace = query
        .workspace_id
        .as_deref()
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(str::to_string)
        .or_else(|| {
            context
                .workspace
                .workspace_id
                .as_deref()
                .map(str::trim)
                .filter(|value| !value.is_empty())
                .map(str::to_string)
        });
    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,
            block_delta_rx,
            stream_delta_rx,
            subscription_workspace,
        }),
        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),
                ));
            }

            // Check block.delta broadcast first（按订阅 workspace 过滤，防跨工作区泄露）
            if let Some(ref mut rx) = state.block_delta_rx {
                match rx.try_recv() {
                    Ok(payload) => {
                        if delta_matches_workspace(
                            &payload,
                            state.subscription_workspace.as_deref(),
                        ) {
                            return Some((Ok(stream_event("block.delta", &payload)), Some(state)));
                        }
                        // 非本工作区：丢弃并继续同一 tick 的后续检查
                    }
                    Err(tokio::sync::broadcast::error::TryRecvError::Empty) => {}
                    Err(tokio::sync::broadcast::error::TryRecvError::Closed) => {
                        state.block_delta_rx = None;
                    }
                    Err(tokio::sync::broadcast::error::TryRecvError::Lagged(_)) => {}
                }
            }

            // Check stream.delta broadcast (push mode: command_committed hints)
            if let Some(ref mut rx) = state.stream_delta_rx {
                match rx.try_recv() {
                    Ok(payload) => {
                        if delta_matches_workspace(
                            &payload,
                            state.subscription_workspace.as_deref(),
                        ) {
                            let hint = build_stream_push_delta_hint(
                                &payload,
                                &state.context.trace.request_id,
                                &state.context.trace.trace_id,
                            );
                            return Some((Ok(stream_event("delta", &hint)), Some(state)));
                        }
                    }
                    Err(tokio::sync::broadcast::error::TryRecvError::Empty) => {}
                    Err(tokio::sync::broadcast::error::TryRecvError::Closed) => {
                        state.stream_delta_rx = None;
                    }
                    Err(tokio::sync::broadcast::error::TryRecvError::Lagged(_)) => {}
                }
            }

            // If push-driven and broadcasts are quiet, polling remains the safety net.
            if state.stream_delta_rx.is_some() {
                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 poll_query = live_poll_query(&state.query);
                let Ok((workspace_id, overview)) =
                    load_stream_overview(state.app_state.config(), &state.context, &poll_query)
                        .await
                else {
                    return None;
                };
                if let Some(change) =
                    resolve_stream_change(&overview, state.current_cursor.as_deref())
                {
                    state.current_cursor = change.cursor.clone();
                    match change.kind {
                        StreamChangeKind::Delta => {
                            let Ok(payload) = build_stream_delta_payload(
                                state.app_state.config(),
                                &state.context,
                                &poll_query,
                                &workspace_id,
                                &overview,
                                change.cursor,
                                change
                                    .delta
                                    .unwrap_or_else(|| serde_json::json!({ "op": "noop" })),
                            )
                            .await
                            else {
                                return None;
                            };
                            return Some((Ok(stream_event("delta", &payload)), Some(state)));
                        }
                        StreamChangeKind::Resync => {
                            let Ok(snapshot_payload) = load_stream_snapshot(
                                state.app_state.config(),
                                &state.context,
                                &poll_query,
                            )
                            .await
                            else {
                                return None;
                            };
                            state.current_cursor =
                                read_stream_cursor_from_payload(&snapshot_payload);
                            return Some((
                                Ok(stream_event(
                                    "resync",
                                    &with_stream_kind(&snapshot_payload, "resync"),
                                )),
                                Some(state),
                            ));
                        }
                    }
                }
                return Some((Ok(stream_event("heartbeat", &json!({}))), Some(state)));
            }

            // Pure polling mode
            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;

                // Check block.delta after poll sleep（同样按 workspace 过滤）
                if let Some(ref mut rx) = state.block_delta_rx {
                    match rx.try_recv() {
                        Ok(payload) => {
                            if delta_matches_workspace(
                                &payload,
                                state.subscription_workspace.as_deref(),
                            ) {
                                return Some((
                                    Ok(stream_event("block.delta", &payload)),
                                    Some(state),
                                ));
                            }
                        }
                        Err(tokio::sync::broadcast::error::TryRecvError::Empty) => {}
                        Err(tokio::sync::broadcast::error::TryRecvError::Closed) => {
                            state.block_delta_rx = None;
                        }
                        Err(tokio::sync::broadcast::error::TryRecvError::Lagged(_)) => {}
                    }
                }

                let poll_query = live_poll_query(&state.query);
                let Ok((workspace_id, overview)) =
                    load_stream_overview(state.app_state.config(), &state.context, &poll_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 Ok(payload) = build_stream_delta_payload(
                            state.app_state.config(),
                            &state.context,
                            &poll_query,
                            &workspace_id,
                            &overview,
                            change.cursor,
                            change
                                .delta
                                .unwrap_or_else(|| serde_json::json!({ "op": "noop" })),
                        )
                        .await
                        else {
                            return None;
                        };
                        return Some((Ok(stream_event("delta", &payload)), Some(state)));
                    }
                    StreamChangeKind::Resync => {
                        let Ok(snapshot_payload) = load_stream_snapshot(
                            state.app_state.config(),
                            &state.context,
                            &poll_query,
                        )
                        .await
                        else {
                            return None;
                        };
                        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 block_delta_rx = state.block_delta_tx.subscribe();
    let stream_delta_rx = state.stream_delta_tx.subscribe();
    let sse = events_with_stream_delta(
        state,
        context,
        query,
        Some(block_delta_rx),
        Some(stream_delta_rx),
    )
    .await?;
    Ok((headers, sse))
}

struct StreamPollState {
    app_state: AppState,
    context: RequestContext,
    query: StreamSnapshotQuery,
    current_cursor: Option<String>,
    polls: u32,
    initial_payload: Value,
    initial_emitted: bool,
    #[allow(dead_code)]
    block_delta_rx: Option<tokio::sync::broadcast::Receiver<Value>>,
    stream_delta_rx: Option<tokio::sync::broadcast::Receiver<Value>>,
    /// 订阅工作区；broadcast 推送仅转发匹配项，避免跨工作区泄露。
    subscription_workspace: Option<String>,
}

/// 仅转发与当前订阅 workspace 一致的 delta（与 ws.rs 同策略）。
fn delta_matches_workspace(delta: &Value, subscription_workspace: Option<&str>) -> bool {
    let Some(expected) = subscription_workspace else {
        // 无订阅 workspace 时不推送带 workspace 的全局 delta（保守）
        return delta
            .get("workspaceId")
            .and_then(|v| v.as_str())
            .map(str::trim)
            .filter(|v| !v.is_empty())
            .is_none();
    };
    match delta
        .get("workspaceId")
        .and_then(|v| v.as_str())
        .map(str::trim)
        .filter(|v| !v.is_empty())
    {
        Some(delta_ws) => delta_ws == expected,
        // 无 workspace 标记的 delta 不转发（避免跨租户噪声）
        None => false,
    }
}

fn live_poll_query(query: &StreamSnapshotQuery) -> StreamSnapshotQuery {
    let mut next = query.clone();
    // Convex bridgeLogs 的 cursor 是“向更旧记录翻页”，不是 live tail 的起点；
    // 实时轮询必须始终查最新窗口，再用 current_cursor 在 Rust 侧比较增量。
    next.cursor = None;
    next
}

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 crate::routes::stream_support::StreamSnapshotQuery;
    use axum::body::{to_bytes, Body};
    use axum::http::{Request, StatusCode};
    use std::time::Duration;
    use tokio::time::timeout;
    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,
            enable_editor_actor: true,
            enable_page_ai_pi_lab: false,
            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(),
            environment: "dev".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\""));
        assert!(text.contains("\"kernel_sidebar_projection\""));
        assert!(text.contains("\"kernel_file_tree_projection\""));
    }

    #[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\""));
    }

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

        assert_eq!(response.status(), StatusCode::OK);
        let body = timeout(
            Duration::from_secs(2),
            to_bytes(response.into_body(), usize::MAX),
        )
        .await
        .expect("push SSE stream should stop after maxPolls")
        .expect("body");
        let text = String::from_utf8(body.to_vec()).expect("utf8");
        assert!(text.contains("event: snapshot") || text.contains("event:snapshot"));
        assert!(text.contains("event: heartbeat") || text.contains("event:heartbeat"));
    }

    #[test]
    fn live_poll_query_drops_bridge_pagination_cursor() {
        let query = StreamSnapshotQuery {
            workspace_id: Some("ws_demo".into()),
            cursor: Some(r#"{"createdAt":"2026-05-12T00:00:00Z","id":"clog_1"}"#.into()),
            poll_ms: Some(250),
            ..StreamSnapshotQuery::default()
        };

        let live_query = super::live_poll_query(&query);

        assert_eq!(live_query.workspace_id, Some("ws_demo".into()));
        assert_eq!(live_query.poll_ms, Some(250));
        assert_eq!(live_query.cursor, None);
    }
}