From d424ea38f0230a24913fdc201614d0abb48681fd Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Mon, 31 Aug 2026 08:45:25 +0800 Subject: [PATCH] perf(ai): batch streaming delta events to the webview --- src-tauri/src/commands/ai.rs | 343 +++++++++++++++++++++++++++++++++-- 1 file changed, 332 insertions(+), 11 deletions(-) diff --git a/src-tauri/src/commands/ai.rs b/src-tauri/src/commands/ai.rs index f0dea97360..7ccf3eda8b 100644 --- a/src-tauri/src/commands/ai.rs +++ b/src-tauri/src/commands/ai.rs @@ -109,10 +109,10 @@ pub async fn ai_stream( ); let cancelled = dbx_core::ai::register_stream(&session_id).await; - let result = dbx_core::ai::stream(&session_id, &request, &cancelled, |chunk| { - let _ = app.emit("ai-stream-chunk", &chunk); - }) - .await; + let batcher = AiStreamChunkBatcher::new(app.clone(), session_id.clone()); + let result = dbx_core::ai::stream(&session_id, &request, &cancelled, |chunk| batcher.handle(chunk)).await; + // Emit tail deltas the interval gate was still holding when the stream ended. + batcher.flush(); dbx_core::ai::unregister_stream(&session_id).await; result @@ -130,6 +130,187 @@ struct AiAgentEventPayload { event: AgentEvent, } +/// Minimum interval between emitted AI delta batches. Providers surface one +/// TextDelta/ReasoningDelta per SSE chunk (30-100+/second) and each emission +/// costs a serde serialization plus a webview IPC dispatch; batching keeps that +/// load off the render path while staying above the ~10Hz rate below which +/// streamed text stops reading as smooth typing. +const AI_DELTA_EMIT_INTERVAL: std::time::Duration = std::time::Duration::from_millis(100); + +/// Delta batches are emitted on the first observation after a quiet period and +/// at most once per interval afterwards. +fn ai_delta_emit_due(last_emit: Option) -> bool { + match last_emit { + None => true, + Some(at) => at.elapsed() >= AI_DELTA_EMIT_INTERVAL, + } +} + +#[derive(Default)] +struct AiAgentDeltaBatch { + text: String, + reasoning: String, + last_emit: Option, +} + +impl AiAgentDeltaBatch { + fn due(&self) -> bool { + ai_delta_emit_due(self.last_emit) + } +} + +/// Coalesces consecutive TextDelta/ReasoningDelta events before they reach the +/// webview. Cloning shares the batch, so provider implementations that clone +/// the `on_event` callback emit through the same buffer. Any non-delta event +/// flushes the pending batch first, preserving event ordering and the +/// exact-replace semantics of WriteSqlConfirmationRequired/ProductionWriteBlocked +/// on the frontend. +struct AiAgentEventBatcher { + app: AppHandle, + session_id: String, + batch: Arc>, +} + +// Manual impl: deriving would add an unnecessary `R: Clone` bound. +impl Clone for AiAgentEventBatcher { + fn clone(&self) -> Self { + Self { app: self.app.clone(), session_id: self.session_id.clone(), batch: self.batch.clone() } + } +} + +impl AiAgentEventBatcher { + fn new(app: AppHandle, session_id: String) -> Self { + Self { app, session_id, batch: Arc::new(std::sync::Mutex::new(AiAgentDeltaBatch::default())) } + } + + fn handle(&self, event: AgentEvent) { + let mut batch = self.lock_batch(); + match event { + AgentEvent::TextDelta { delta } => { + batch.text.push_str(&delta); + if batch.due() { + self.flush_locked(&mut batch); + } + } + AgentEvent::ReasoningDelta { delta } => { + batch.reasoning.push_str(&delta); + if batch.due() { + self.flush_locked(&mut batch); + } + } + event => { + self.flush_locked(&mut batch); + self.emit_event(event); + } + } + } + + /// Emits whatever deltas the interval gate is still holding. + fn flush(&self) { + let mut batch = self.lock_batch(); + self.flush_locked(&mut batch); + } + + fn lock_batch(&self) -> std::sync::MutexGuard<'_, AiAgentDeltaBatch> { + self.batch.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) + } + + fn flush_locked(&self, batch: &mut AiAgentDeltaBatch) { + if !batch.text.is_empty() { + self.emit_event(AgentEvent::TextDelta { delta: std::mem::take(&mut batch.text) }); + } + if !batch.reasoning.is_empty() { + self.emit_event(AgentEvent::ReasoningDelta { delta: std::mem::take(&mut batch.reasoning) }); + } + batch.last_emit = Some(std::time::Instant::now()); + } + + fn emit_event(&self, event: AgentEvent) { + let payload = AiAgentEventPayload { session_id: self.session_id.clone(), event }; + let _ = self.app.emit("ai-agent-event", &payload); + } +} + +#[derive(Default)] +struct AiStreamChunkBatch { + delta: String, + reasoning: Option, + last_emit: Option, +} + +impl AiStreamChunkBatch { + fn due(&self) -> bool { + ai_delta_emit_due(self.last_emit) + } +} + +/// Same delta coalescing as [`AiAgentEventBatcher`] for the legacy +/// `ai-stream-chunk` completion stream: concatenate deltas between emissions, +/// forward the terminal `done` chunk immediately after flushing. +struct AiStreamChunkBatcher { + app: AppHandle, + session_id: String, + batch: Arc>, +} + +impl Clone for AiStreamChunkBatcher { + fn clone(&self) -> Self { + Self { app: self.app.clone(), session_id: self.session_id.clone(), batch: self.batch.clone() } + } +} + +impl AiStreamChunkBatcher { + fn new(app: AppHandle, session_id: String) -> Self { + Self { app, session_id, batch: Arc::new(std::sync::Mutex::new(AiStreamChunkBatch::default())) } + } + + fn handle(&self, mut chunk: AiStreamChunk) { + let mut batch = self.lock_batch(); + if chunk.done { + self.flush_locked(&mut batch); + self.emit_chunk(chunk); + return; + } + if !chunk.delta.is_empty() { + batch.delta.push_str(&chunk.delta); + } + if let Some(reasoning) = chunk.reasoning_delta.take() { + if !reasoning.is_empty() { + batch.reasoning.get_or_insert_with(String::new).push_str(&reasoning); + } + } + if batch.due() { + self.flush_locked(&mut batch); + } + } + + /// Emits whatever deltas the interval gate is still holding. + fn flush(&self) { + let mut batch = self.lock_batch(); + self.flush_locked(&mut batch); + } + + fn lock_batch(&self) -> std::sync::MutexGuard<'_, AiStreamChunkBatch> { + self.batch.lock().unwrap_or_else(|poisoned| poisoned.into_inner()) + } + + fn flush_locked(&self, batch: &mut AiStreamChunkBatch) { + if !batch.delta.is_empty() || batch.reasoning.is_some() { + self.emit_chunk(AiStreamChunk { + session_id: self.session_id.clone(), + delta: std::mem::take(&mut batch.delta), + reasoning_delta: batch.reasoning.take(), + done: false, + }); + } + batch.last_emit = Some(std::time::Instant::now()); + } + + fn emit_chunk(&self, chunk: AiStreamChunk) { + let _ = self.app.emit("ai-stream-chunk", &chunk); + } +} + #[tauri::command] pub async fn ai_cancel_stream(session_id: String) -> Result { Ok(dbx_core::ai::cancel_stream(&session_id).await) @@ -214,18 +395,15 @@ pub async fn ai_agent_stream( }; let is_agent_mode = mode.as_deref() == Some("agent"); + let emitter = AiAgentEventBatcher::new(app.clone(), session_id.clone()); let result = run_agent_loop( &request.config, &request.system_prompt, &request.messages, &agent_ctx, { - let app = app.clone(); - let event_session_id = session_id.clone(); - move |event: AgentEvent| { - let payload = AiAgentEventPayload { session_id: event_session_id.clone(), event }; - let _ = app.emit("ai-agent-event", &payload); - } + let emitter = emitter.clone(); + move |event: AgentEvent| emitter.handle(event) }, &cancelled, request.max_tokens, @@ -233,6 +411,9 @@ pub async fn ai_agent_stream( is_agent_mode, ) .await; + // Emit tail deltas the interval gate was still holding when the loop ended + // (e.g. an error return that produced no terminal AgentEvent). + emitter.flush(); dbx_core::ai::unregister_stream(&session_id).await; result @@ -310,12 +491,152 @@ mod tests { use std::sync::atomic::{AtomicU32, Ordering}; use std::sync::Arc; - use tauri::Manager; + use tauri::{Listener, Manager}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use super::super::connection::AppState; use dbx_core::ai::{AiApiStyle, AiAuthMethod, AiConfig, AiProvider, AiReasoningLevel}; + /// Captures the JSON payload of every `event_name` emission on a mock app. + fn captured_events( + event_name: &'static str, + ) -> (tauri::AppHandle, Arc>>) { + let app = tauri::test::mock_app(); + let handle = app.handle().clone(); + let received: Arc>> = Arc::default(); + let sink = received.clone(); + handle.listen(event_name, move |event| { + if let Ok(value) = serde_json::from_str::(event.payload()) { + sink.lock().unwrap().push(value); + } + }); + (handle, received) + } + + #[test] + fn ai_agent_event_batcher_coalesces_deltas_between_intervals() { + use super::AiAgentEventBatcher; + use dbx_core::agent_events::AgentEvent; + + let (handle, received) = captured_events("ai-agent-event"); + let batcher = AiAgentEventBatcher::new(handle, "session-1".to_string()); + // Leading delta flushes immediately. + batcher.handle(AgentEvent::TextDelta { delta: "a".to_string() }); + // Follow-up deltas within the interval are held … + batcher.handle(AgentEvent::TextDelta { delta: "b".to_string() }); + batcher.handle(AgentEvent::ReasoningDelta { delta: "r1".to_string() }); + // … and released as one batch each before a non-delta event, so event + // ordering (and the exact-replace semantics of confirmation events on + // the frontend) is preserved. + batcher.handle(AgentEvent::TurnStart { turn: 1 }); + + let observed: Vec<(String, String)> = received + .lock() + .unwrap() + .iter() + .map(|value| { + let kind = value.get("type").and_then(|v| v.as_str()).unwrap_or_default().to_string(); + let delta = value.get("delta").and_then(|v| v.as_str()).unwrap_or_default().to_string(); + (kind, delta) + }) + .collect(); + assert_eq!( + observed, + vec![ + ("text_delta".to_string(), "a".to_string()), + ("text_delta".to_string(), "b".to_string()), + ("reasoning_delta".to_string(), "r1".to_string()), + ("turn_start".to_string(), String::new()), + ] + ); + + // Tail deltas the interval gate is still holding flush at stream end. + received.lock().unwrap().clear(); + batcher.handle(AgentEvent::TextDelta { delta: "c".to_string() }); + batcher.flush(); + let observed: Vec<(String, String)> = received + .lock() + .unwrap() + .iter() + .map(|value| { + let kind = value.get("type").and_then(|v| v.as_str()).unwrap_or_default().to_string(); + let delta = value.get("delta").and_then(|v| v.as_str()).unwrap_or_default().to_string(); + (kind, delta) + }) + .collect(); + assert_eq!(observed, vec![("text_delta".to_string(), "c".to_string())]); + } + + #[test] + fn ai_agent_event_batcher_merges_held_deltas_after_interval() { + use super::AiAgentEventBatcher; + use dbx_core::agent_events::AgentEvent; + + let (handle, received) = captured_events("ai-agent-event"); + let batcher = AiAgentEventBatcher::new(handle, "session-1".to_string()); + batcher.handle(AgentEvent::TextDelta { delta: "a".to_string() }); + batcher.handle(AgentEvent::TextDelta { delta: "b".to_string() }); + std::thread::sleep(super::AI_DELTA_EMIT_INTERVAL + std::time::Duration::from_millis(20)); + batcher.handle(AgentEvent::TextDelta { delta: "c".to_string() }); + + let deltas: Vec = received + .lock() + .unwrap() + .iter() + .filter(|value| value.get("type").and_then(|v| v.as_str()) == Some("text_delta")) + .filter_map(|value| value.get("delta").and_then(|v| v.as_str()).map(ToString::to_string)) + .collect(); + assert_eq!(deltas, vec!["a".to_string(), "bc".to_string()]); + } + + #[test] + fn ai_stream_chunk_batcher_merges_deltas_and_forwards_done() { + use super::AiStreamChunkBatcher; + + let (handle, received) = captured_events("ai-stream-chunk"); + let batcher = AiStreamChunkBatcher::new(handle, "session-2".to_string()); + batcher.handle(super::AiStreamChunk { + session_id: "session-2".to_string(), + delta: "a".to_string(), + reasoning_delta: None, + done: false, + }); + // Held within the interval … + batcher.handle(super::AiStreamChunk { + session_id: "session-2".to_string(), + delta: "b".to_string(), + reasoning_delta: Some("r".to_string()), + done: false, + }); + // … and flushed in order before the terminal chunk is forwarded. + batcher.handle(super::AiStreamChunk { + session_id: "session-2".to_string(), + delta: String::new(), + reasoning_delta: None, + done: true, + }); + + let observed: Vec<(String, Option, bool)> = received + .lock() + .unwrap() + .iter() + .map(|value| { + let delta = value.get("delta").and_then(|v| v.as_str()).unwrap_or_default().to_string(); + let reasoning = value.get("reasoning_delta").and_then(|v| v.as_str()).map(ToString::to_string); + let done = value.get("done").and_then(|v| v.as_bool()).unwrap_or(false); + (delta, reasoning, done) + }) + .collect(); + assert_eq!( + observed, + vec![ + ("a".to_string(), None, false), + ("b".to_string(), Some("r".to_string()), false), + (String::new(), None, true), + ] + ); + } + /// Spawn a TCP server that returns 429 and counts connections. async fn counting_429_server() -> (String, Arc, tokio::task::JoinHandle<()>) { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();