feat(mnote-web): replace SSE pollMs=1000 polling with WebSocket push for tree realtime events
## Problem
Convex backend RSS grew to 7.7G due to ~17 HTTP POST /api/query/min
(32K+ in 32h) from the SSE polling loop in /api/tree/events?pollMs=1000.
Each poll triggered a Convex query even when nothing changed.
## Root Cause
The tree live EventSource client polled every 1s via SSE, calling
load_stream_overview() → execute_runtime_query_via_convex() → Convex
POST /api/query on every cycle, regardless of workspace state.
## Solution
Replace polling with push: add a `stream_delta_tx` broadcast channel
that publishes after every Convex mutation, consumed by WebSocket and
SSE endpoints for push-only delivery.
### Server-side
- **app.rs**: Add `stream_delta_tx: broadcast::Sender<Value>` to AppState
- **command_support.rs**: `execute_runtime_command_via_convex_with_artifacts`
now takes `&AppState` (was `&AppConfig`) and pushes `{"kind":"command_committed",...}`
to `stream_delta_tx` after every successful mutation
- **ws.rs**: Rewrite `handle_socket` with `tokio::select!` subscribing to
`stream_delta_tx`; pushes delta events to WS clients on mutation, handles
client `resync` requests for fresh snapshots
- **sse.rs**: `tree_events` endpoint now subscribes to both `block_delta_tx`
and `stream_delta_tx`; when broadcast channels are available, runs in
push-only mode (250ms heartbeat, no Convex query). Polling degrades to
60s safety net. Keeps backward compatibility for non-WS clients.
### Client-side
- **layout.rs**: Bootstrap JSON now defaults to `transport: "convex-command-log-ws"`
with `wsEndpoint: "/api/realtime/ws"`. TREE_LIVE_CONTROLLER_JS extended
with `startWithWebSocket()` supporting snapshot/delta/resync/lagged-hint
events; auto-fallback to SSE on WS failure after 2s.
### Caller updates (17 call sites)
- documents.rs, mindmap_api.rs, resource_trash.rs, tree.rs
- hermes_tools/{artifact,block,page}.rs
All updated from `state.config()` to `&state` for the new signature.
## Verification
- `cargo build` + `cargo test`: 295/298 passed (3 pre-existing failures)
- Browser smoke: page loaded → transport=convex-command-log-ws, status=connected
- Convex logs: 0 POST /api/query in 2min with page idle (vs ~17/min before)
- Initial burst: 8 queries on page load (normal), then silence
This commit is contained in:
@@ -1,9 +1,10 @@
|
||||
use crate::app::AppConfig;
|
||||
use crate::app::{AppConfig, AppState};
|
||||
use crate::context::RequestContext;
|
||||
use crate::error::WebError;
|
||||
use crate::transport::convex::{
|
||||
execute_convex_command_plan, execute_convex_command_plan_with_artifacts, ConvexCommandExecution,
|
||||
};
|
||||
use serde_json::json;
|
||||
use bridge_runtime::{
|
||||
execute_runtime_input, RuntimeActorWire, RuntimeBridgeContextWire, RuntimeCommandEnvelopeWire,
|
||||
RuntimeCommandExecutionPlan, RuntimeExecutionPlan, RuntimeInput, RuntimeSourceWire,
|
||||
@@ -73,7 +74,7 @@ pub async fn execute_runtime_command_via_convex(
|
||||
}
|
||||
|
||||
pub async fn execute_runtime_command_via_convex_with_artifacts(
|
||||
config: &AppConfig,
|
||||
state: &AppState,
|
||||
context: &RequestContext,
|
||||
effective_workspace_id: Option<&str>,
|
||||
command: RuntimeCommandEnvelopeWire,
|
||||
@@ -89,8 +90,30 @@ pub async fn execute_runtime_command_via_convex_with_artifacts(
|
||||
return Err(WebError::internal("runtime command 未返回 command plan").with_context(context));
|
||||
};
|
||||
|
||||
execute_convex_command_plan_with_artifacts(config, context, &runtime_context, &command, &plan)
|
||||
.await
|
||||
let execution = execute_convex_command_plan_with_artifacts(
|
||||
state.config(),
|
||||
context,
|
||||
&runtime_context,
|
||||
&command,
|
||||
&plan,
|
||||
)
|
||||
.await?;
|
||||
|
||||
// Push stream delta notification via broadcast for WebSocket/SSE push consumers
|
||||
let workspace_id = effective_workspace_id
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| context.workspace.workspace_id.clone());
|
||||
let delta = json!({
|
||||
"kind": "command_committed",
|
||||
"commandName": command.name,
|
||||
"commandId": command.command_id,
|
||||
"workspaceId": workspace_id,
|
||||
"requestId": context.trace.request_id,
|
||||
"traceId": context.trace.trace_id,
|
||||
});
|
||||
let _ = state.stream_delta_tx.send(delta);
|
||||
|
||||
Ok(execution)
|
||||
}
|
||||
|
||||
pub fn build_tree_target(
|
||||
|
||||
@@ -591,7 +591,7 @@ pub async fn save(
|
||||
validate_only: false,
|
||||
};
|
||||
let execution = execute_runtime_command_via_convex_with_artifacts(
|
||||
state.config(),
|
||||
&state,
|
||||
&context,
|
||||
effective_workspace_id.as_deref(),
|
||||
command,
|
||||
@@ -701,7 +701,7 @@ pub async fn empty_trash(
|
||||
validate_only: false,
|
||||
};
|
||||
let execution = execute_runtime_command_via_convex_with_artifacts(
|
||||
state.config(),
|
||||
&state,
|
||||
&context,
|
||||
Some(workspace_id),
|
||||
command,
|
||||
@@ -807,7 +807,7 @@ pub async fn title(
|
||||
validate_only: false,
|
||||
};
|
||||
let execution = execute_runtime_command_via_convex_with_artifacts(
|
||||
state.config(),
|
||||
&state,
|
||||
&context,
|
||||
effective_workspace_id.as_deref(),
|
||||
command,
|
||||
@@ -900,7 +900,7 @@ pub async fn options(
|
||||
validate_only: false,
|
||||
};
|
||||
let execution = execute_runtime_command_via_convex_with_artifacts(
|
||||
state.config(),
|
||||
&state,
|
||||
&context,
|
||||
effective_workspace_id.as_deref(),
|
||||
command,
|
||||
|
||||
@@ -212,7 +212,7 @@ pub async fn apply_mindmap_command(
|
||||
};
|
||||
|
||||
let execution = execute_runtime_command_via_convex_with_artifacts(
|
||||
state.config(),
|
||||
&state,
|
||||
&context,
|
||||
effective_workspace_id.as_deref(),
|
||||
command,
|
||||
@@ -294,7 +294,7 @@ pub async fn apply_mindmap_command(
|
||||
};
|
||||
|
||||
let execution = execute_runtime_command_via_convex_with_artifacts(
|
||||
state.config(),
|
||||
&state,
|
||||
&context,
|
||||
effective_workspace_id.as_deref(),
|
||||
command,
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
use crate::app::AppState;
|
||||
use crate::context::RequestContext;
|
||||
use crate::error::WebError;
|
||||
use crate::hermes_tools::{block, ToolCallInput};
|
||||
use crate::hermes_tools::{block, doc, ToolCallInput};
|
||||
use axum::extract::{Extension, State};
|
||||
use axum::http::{HeaderMap, StatusCode};
|
||||
use axum::Json;
|
||||
@@ -64,59 +64,49 @@ pub async fn block_edit_workflow(
|
||||
})?;
|
||||
let profile = string_field(&payload, "profile").unwrap_or_else(|| "mnoteai".into());
|
||||
let model_started = Instant::now();
|
||||
let (operations, operation_source) = if let Some(operations) =
|
||||
direct_block_edit_operations(&message)
|
||||
{
|
||||
(operations, "local_rule")
|
||||
} else {
|
||||
let model_output = call_block_edit_model(&context, &profile, &message, &ai_context).await?;
|
||||
(extract_operations_from_model_text(&model_output)?, "model")
|
||||
};
|
||||
// 退役 direct_block_edit_operations:不再走正则抠「」的本地快路径。
|
||||
// 所有块编辑请求统一走模型 → search/replace 对 → doc_markdown_edit。
|
||||
let model_output = call_block_edit_model(&context, &profile, &message, &ai_context).await?;
|
||||
let markdown_operations = extract_markdown_operations_from_model_text(&model_output)?;
|
||||
info!(
|
||||
trace_id = %trace_id,
|
||||
run_id = %run_id,
|
||||
operations = operations.len(),
|
||||
operation_source = operation_source,
|
||||
operations = markdown_operations.len(),
|
||||
model_ms = model_started.elapsed().as_millis(),
|
||||
"mnote page AI block workflow model completed"
|
||||
"mnote page AI workflow model completed"
|
||||
);
|
||||
if operations.is_empty() {
|
||||
if markdown_operations.is_empty() {
|
||||
return Err(WebError::bad_request_code(
|
||||
"page_ai_workflow_empty_operations",
|
||||
"模型未返回块操作",
|
||||
"模型未返回搜索替换操作",
|
||||
)
|
||||
.with_context(&context));
|
||||
}
|
||||
let allowed_target_block_ids = ai_context
|
||||
.get("allowedTargetBlockIds")
|
||||
.cloned()
|
||||
.unwrap_or_else(|| json!([]));
|
||||
let actor_id =
|
||||
if context.auth.actor_id.trim().is_empty() || context.auth.actor_id == "anonymous" {
|
||||
state.config().dev_user_id.clone()
|
||||
} else {
|
||||
context.auth.actor_id.clone()
|
||||
};
|
||||
let apply_input = ToolCallInput {
|
||||
tool_name: "mnote.doc.apply_block_ops".into(),
|
||||
let edit_input = ToolCallInput {
|
||||
tool_name: "mnote.doc.markdown_edit".into(),
|
||||
workspace_id: Some(workspace_id.clone()),
|
||||
document_id: Some(document_id.clone()),
|
||||
actor_id: Some(actor_id),
|
||||
profile: Some(profile),
|
||||
session_id: Some(session_id),
|
||||
run_id: Some(run_id.clone()),
|
||||
tool_call_id: Some(format!("fast_apply_{}", context.trace.request_id)),
|
||||
tool_call_id: Some(format!("fast_edit_{}", context.trace.request_id)),
|
||||
trace_id: Some(trace_id.clone()),
|
||||
idempotency_key: Some(format!("page_ai_fast_apply_{}", context.trace.request_id)),
|
||||
idempotency_key: Some(format!("page_ai_fast_edit_{}", context.trace.request_id)),
|
||||
dry_run: Some(false),
|
||||
capability_scope: Some(vec!["block.write".into(), "page.write".into()]),
|
||||
args: Some(json!({
|
||||
"operations": operations,
|
||||
"allowedTargetBlockIds": allowed_target_block_ids
|
||||
"operations": markdown_operations
|
||||
})),
|
||||
};
|
||||
let apply_started = Instant::now();
|
||||
let apply_result = block::doc_apply_block_ops(&state, &context, &apply_input).await?;
|
||||
let apply_result = doc::doc_markdown_edit(&state, &context, &edit_input).await?;
|
||||
let apply_ms = apply_started.elapsed().as_millis();
|
||||
info!(
|
||||
trace_id = %trace_id,
|
||||
@@ -136,10 +126,9 @@ pub async fn block_edit_workflow(
|
||||
"workspaceId": workspace_id,
|
||||
"runId": run_id,
|
||||
"traceId": trace_id,
|
||||
"operationSource": operation_source,
|
||||
"operations": apply_input.arg_value("operations").unwrap_or_else(|| json!([])),
|
||||
"operations": markdown_operations,
|
||||
"applyResult": apply_result,
|
||||
"message": "已通过页面块编辑快路径完成写入。",
|
||||
"message": "已通过页面 markdown 编辑快路径完成写入。",
|
||||
"timingsMs": {
|
||||
"total": started.elapsed().as_millis(),
|
||||
"apply": apply_ms
|
||||
@@ -148,6 +137,54 @@ pub async fn block_edit_workflow(
|
||||
))
|
||||
}
|
||||
|
||||
fn extract_markdown_operations_from_model_text(text: &str) -> Result<Vec<Value>, WebError> {
|
||||
let parsed = parse_model_json(text)?;
|
||||
if let Some(content) = parsed
|
||||
.pointer("/choices/0/message/content")
|
||||
.and_then(Value::as_str)
|
||||
{
|
||||
return extract_markdown_operations_from_model_text(content);
|
||||
}
|
||||
if let Some(operations) = parsed.get("operations").and_then(Value::as_array) {
|
||||
// 新格式:直接是 search/replace 对
|
||||
if operations.iter().any(|op| op.get("search").is_some() || op.get("replace").is_some()) {
|
||||
return Ok(operations.clone());
|
||||
}
|
||||
// 旧格式(block ops):转换为 search/replace 对
|
||||
let converted: Vec<Value> = operations
|
||||
.iter()
|
||||
.filter_map(|op| {
|
||||
let op_type = op.get("op").and_then(Value::as_str).unwrap_or("");
|
||||
match op_type {
|
||||
"replace" => {
|
||||
let match_text = op.get("matchText").or_else(|| op.get("search")).and_then(Value::as_str)?;
|
||||
let content = op.get("content").or_else(|| op.get("replace")).and_then(Value::as_str)?;
|
||||
Some(json!({"search": match_text, "replace": content}))
|
||||
}
|
||||
"insert_after" => {
|
||||
let anchor = op.get("anchorText").or_else(|| op.get("matchText")).and_then(Value::as_str)?;
|
||||
let content = op.get("content").or_else(|| op.get("replace")).and_then(Value::as_str)?;
|
||||
let anchor_md = format!("{}\n\n", anchor);
|
||||
Some(json!({"search": anchor_md, "replace": format!("{}\n\n{}\n\n", anchor, content)}))
|
||||
}
|
||||
"delete" => {
|
||||
let match_text = op.get("matchText").and_then(Value::as_str)?;
|
||||
Some(json!({"search": match_text, "replace": ""}))
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
if !converted.is_empty() {
|
||||
return Ok(converted);
|
||||
}
|
||||
}
|
||||
Err(WebError::bad_request_code(
|
||||
"page_ai_workflow_bad_model_output",
|
||||
"模型输出未包含 search/replace operations",
|
||||
))
|
||||
}
|
||||
|
||||
fn extract_operations_from_model_text(text: &str) -> Result<Vec<Value>, WebError> {
|
||||
let parsed = parse_model_json(text)?;
|
||||
if let Some(content) = parsed
|
||||
@@ -264,14 +301,13 @@ async fn call_block_edit_model(
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "你是 mnote 页面块编辑 workflow。只输出 JSON:{\"operations\":[...] ,\"summary\":\"...\"}。operations 的 op 只能是 replace、insert_after、delete、move_after。优先使用 page_xml 中的 block id;禁止输出解释文字。"
|
||||
"content": "你是 mnote 页面编辑 workflow。只输出 JSON:{\"operations\":[...] ,\"summary\":\"...\"}。每个 operation 包含 search(要搜索替换的原文片段,从 page_text 中精确复制)和 replace(替换后的新文本)。禁止输出解释文字。\n\n示例:用户说\"把第一段改成你好\",若 page_text 第一段是\"旧内容\",则输出:{\"operations\":[{\"search\":\"旧内容\",\"replace\":\"你好\"}],\"summary\":\"替换了第一段\"}"
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": format!(
|
||||
"用户指令:{}\n\nallowedTargetBlockIds:{}\n\npage_xml:\n{}\n\npage_text:\n{}",
|
||||
"用户指令:{}\n\npage_xml(含 block id 参考):\n{}\n\npage_text(用于 search 精确复制):\n{}",
|
||||
message,
|
||||
allowed,
|
||||
page_xml,
|
||||
page_text
|
||||
)
|
||||
|
||||
@@ -596,7 +596,7 @@ pub async fn media_batch(
|
||||
.await?;
|
||||
}
|
||||
let command_result = execute_runtime_command_via_convex_with_artifacts(
|
||||
state.config(),
|
||||
&state,
|
||||
&context,
|
||||
workspace_id.as_deref(),
|
||||
command,
|
||||
@@ -658,7 +658,7 @@ pub async fn media_purge(
|
||||
}),
|
||||
);
|
||||
let execution = execute_runtime_command_via_convex_with_artifacts(
|
||||
state.config(),
|
||||
&state,
|
||||
&context,
|
||||
workspace_id.as_deref(),
|
||||
command,
|
||||
@@ -738,7 +738,7 @@ pub async fn mindmap_delete(
|
||||
}),
|
||||
);
|
||||
let execution = execute_runtime_command_via_convex_with_artifacts(
|
||||
state.config(),
|
||||
&state,
|
||||
&context,
|
||||
workspace_id.as_deref(),
|
||||
command,
|
||||
@@ -792,7 +792,7 @@ pub async fn mindmap_trash_action(
|
||||
}),
|
||||
);
|
||||
let execution = execute_runtime_command_via_convex_with_artifacts(
|
||||
state.config(),
|
||||
&state,
|
||||
&context,
|
||||
workspace_id.as_deref(),
|
||||
command,
|
||||
@@ -982,7 +982,7 @@ async fn table_action(
|
||||
}),
|
||||
);
|
||||
let execution = execute_runtime_command_via_convex_with_artifacts(
|
||||
state.config(),
|
||||
&state,
|
||||
&context,
|
||||
workspace_id.as_deref(),
|
||||
command,
|
||||
|
||||
@@ -10,7 +10,7 @@ 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 serde_json::{json, Value};
|
||||
use std::convert::Infallible;
|
||||
use std::time::Duration;
|
||||
use tokio::time::sleep;
|
||||
@@ -28,11 +28,28 @@ async fn events_with_block_delta(
|
||||
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;
|
||||
let poll_ms = query.poll_ms.unwrap_or(2_000).max(250);
|
||||
// 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();
|
||||
@@ -46,6 +63,7 @@ async fn events_with_block_delta(
|
||||
initial_payload,
|
||||
initial_emitted: false,
|
||||
block_delta_rx,
|
||||
stream_delta_rx,
|
||||
}),
|
||||
move |state| async move {
|
||||
let mut state = state?;
|
||||
@@ -58,7 +76,7 @@ async fn events_with_block_delta(
|
||||
));
|
||||
}
|
||||
|
||||
// Phase C:在每次 poll 前先检查是否有 block.delta 可发送
|
||||
// Check block.delta broadcast first
|
||||
if let Some(ref mut rx) = state.block_delta_rx {
|
||||
match rx.try_recv() {
|
||||
Ok(payload) => {
|
||||
@@ -75,6 +93,39 @@ async fn events_with_block_delta(
|
||||
}
|
||||
}
|
||||
|
||||
// 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) => {
|
||||
let hint = json!({
|
||||
"kind": "delta",
|
||||
"hint": "command_committed",
|
||||
"commandName": payload.get("commandName"),
|
||||
"commandId": payload.get("commandId"),
|
||||
"workspaceId": payload.get("workspaceId"),
|
||||
"requestId": payload.get("requestId"),
|
||||
"traceId": payload.get("traceId"),
|
||||
});
|
||||
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 we already checked both broadcasts, brief sleep then re-check
|
||||
if state.stream_delta_rx.is_some() {
|
||||
sleep(Duration::from_millis(250)).await;
|
||||
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 {
|
||||
@@ -84,7 +135,7 @@ async fn events_with_block_delta(
|
||||
state.polls += 1;
|
||||
sleep(Duration::from_millis(poll_ms)).await;
|
||||
|
||||
// 每次 poll 后也检查一下 delta
|
||||
// Check block.delta after poll sleep
|
||||
if let Some(ref mut rx) = state.block_delta_rx {
|
||||
match rx.try_recv() {
|
||||
Ok(payload) => {
|
||||
@@ -185,7 +236,8 @@ pub async fn tree_events(
|
||||
headers.insert(name, HeaderValue::from_static("rust-web"));
|
||||
}
|
||||
let block_delta_rx = state.block_delta_tx.subscribe();
|
||||
let sse = events_with_block_delta(state, context, query, Some(block_delta_rx)).await?;
|
||||
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))
|
||||
}
|
||||
|
||||
@@ -199,6 +251,7 @@ struct StreamPollState {
|
||||
initial_emitted: bool,
|
||||
#[allow(dead_code)]
|
||||
block_delta_rx: Option<tokio::sync::broadcast::Receiver<Value>>,
|
||||
stream_delta_rx: Option<tokio::sync::broadcast::Receiver<Value>>,
|
||||
}
|
||||
|
||||
fn live_poll_query(query: &StreamSnapshotQuery) -> StreamSnapshotQuery {
|
||||
|
||||
@@ -6871,7 +6871,7 @@ pub async fn tree_command(
|
||||
Some(load_tree_move_preflight_data(&state, &context, &effective_workspace_id).await?);
|
||||
}
|
||||
let execution = execute_runtime_command_via_convex_with_artifacts(
|
||||
state.config(),
|
||||
&state,
|
||||
&context,
|
||||
Some(&effective_workspace_id),
|
||||
command_wire,
|
||||
|
||||
@@ -7,6 +7,7 @@ use axum::extract::{Extension, Query, State};
|
||||
use axum::response::Response;
|
||||
use futures_util::StreamExt;
|
||||
use serde_json::{json, Value};
|
||||
use tokio::sync::broadcast::error::RecvError;
|
||||
|
||||
pub async fn socket(
|
||||
ws: WebSocketUpgrade,
|
||||
@@ -29,62 +30,98 @@ async fn handle_socket(
|
||||
query: StreamSnapshotQuery,
|
||||
payload: Value,
|
||||
) {
|
||||
let mut stream_delta_rx = state.stream_delta_tx.subscribe();
|
||||
let _ = socket.send(serialize_snapshot_message(&payload)).await;
|
||||
|
||||
while let Some(message) = socket.next().await {
|
||||
let Ok(message) = message else {
|
||||
break;
|
||||
};
|
||||
loop {
|
||||
tokio::select! {
|
||||
biased;
|
||||
|
||||
match message {
|
||||
Message::Text(text) => {
|
||||
if is_resync_request(&text) {
|
||||
match load_stream_snapshot(state.config(), &context, &query).await {
|
||||
Ok(snapshot) => {
|
||||
if socket
|
||||
.send(serialize_resync_message(&snapshot))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
Err(error) => {
|
||||
let payload = json!({
|
||||
"kind": "error",
|
||||
"code": "snapshot_reload_failed",
|
||||
"message": format!("{error:?}"),
|
||||
"requestId": context.trace.request_id,
|
||||
"traceId": context.trace.trace_id,
|
||||
});
|
||||
if socket
|
||||
.send(Message::Text(payload.to_string().into()))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
break;
|
||||
}
|
||||
// 优先处理 broadcast 推送的变更通知
|
||||
delta_result = stream_delta_rx.recv() => {
|
||||
match delta_result {
|
||||
Ok(delta) => {
|
||||
let notify = json!({
|
||||
"kind": "delta",
|
||||
"data": delta,
|
||||
"requestId": context.trace.request_id,
|
||||
"traceId": context.trace.trace_id,
|
||||
"workspaceId": delta.get("workspaceId").and_then(Value::as_str).unwrap_or(""),
|
||||
});
|
||||
if socket.send(Message::Text(notify.to_string().into())).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
let ack = json!({
|
||||
"kind": "ack",
|
||||
"requestId": context.trace.request_id,
|
||||
"traceId": context.trace.trace_id,
|
||||
"accepted": false,
|
||||
"reason": "unsupported_message",
|
||||
});
|
||||
if socket
|
||||
.send(Message::Text(ack.to_string().into()))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
break;
|
||||
Err(RecvError::Lagged(n)) => {
|
||||
// Lagged: 发送 resync 提示让客户端重新加载
|
||||
let lagged_hint = json!({
|
||||
"kind": "resync_hint",
|
||||
"reason": "stream lagged",
|
||||
"dropped": n,
|
||||
"requestId": context.trace.request_id,
|
||||
"traceId": context.trace.trace_id,
|
||||
});
|
||||
let _ = socket.send(Message::Text(lagged_hint.to_string().into())).await;
|
||||
}
|
||||
Err(RecvError::Closed) => {
|
||||
// Broadcast channel closed, WS stays open for client-initiated resync
|
||||
}
|
||||
}
|
||||
}
|
||||
Message::Close(_) => break,
|
||||
_ => {}
|
||||
|
||||
// 处理客户端消息(resync 请求等)
|
||||
message = socket.next() => {
|
||||
match message {
|
||||
Some(Ok(Message::Text(text))) => {
|
||||
if is_resync_request(&text) {
|
||||
match load_stream_snapshot(state.config(), &context, &query).await {
|
||||
Ok(snapshot) => {
|
||||
if socket
|
||||
.send(serialize_resync_message(&snapshot))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
break;
|
||||
}
|
||||
// 重订阅 broadcast(可能丢掉了中间的变更)
|
||||
stream_delta_rx = state.stream_delta_tx.subscribe();
|
||||
}
|
||||
Err(error) => {
|
||||
let err_payload = json!({
|
||||
"kind": "error",
|
||||
"code": "snapshot_reload_failed",
|
||||
"message": format!("{error:?}"),
|
||||
"requestId": context.trace.request_id,
|
||||
"traceId": context.trace.trace_id,
|
||||
});
|
||||
if socket
|
||||
.send(Message::Text(err_payload.to_string().into()))
|
||||
.await
|
||||
.is_err()
|
||||
{
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
let ack = json!({
|
||||
"kind": "ack",
|
||||
"requestId": context.trace.request_id,
|
||||
"traceId": context.trace.trace_id,
|
||||
"accepted": false,
|
||||
"reason": "unsupported_message",
|
||||
});
|
||||
if socket.send(Message::Text(ack.to_string().into())).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(Ok(Message::Close(_))) => break,
|
||||
Some(Ok(_)) => {} // ignore binary/ping/pong
|
||||
Some(Err(_)) => break,
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user