Batch completion events to reduce notify() calls during streaming (#46802)

## Problem

Profiling showed that during agent streaming, `thread.rs:1393:23` was
appearing constantly as a hotspot. The issue was that every single token
from the model triggered:
1. `this.update(cx, ...)` 
2. `handle_text_event()` (or thinking/redacted_thinking)
3. `cx.notify()`
4. UI re-render

This created significant foreground thread pressure, contributing to
~500ms delays visible in the profiler.

## Solution

Batch all immediately-available events using `now_or_never()` and
process them in a single `update()` call with one `notify()` at the end.

This approach is deterministic - it processes exactly what's available
right now, adapting naturally to network speed:
- When tokens arrive slowly, you get one at a time
- When they arrive in bursts, they batch automatically

## Changes

- Remove `cx.notify()` from `handle_text_event`,
`handle_thinking_event`, and `handle_redacted_thinking_event`
- Batch events in the streaming loop using `now_or_never()`
- Call `cx.notify()` once per batch instead of per-event
- Keep `cx.notify()` in `handle_tool_use_event` for immediate tool
feedback

Release Notes:
- Improve streaming tool call performance by batching UI updates.
This commit is contained in:
Richard Feldman authored and GitHub committed 2026-01-14 12:14:28 -05:00
1 parent cb51571b5b
commit e5ca2f0aae
1 file changed
+52 -34
+52 -34
View File
@@ -1471,8 +1471,8 @@ impl Thread {
let mut tool_results = FuturesUnordered::new();
let mut cancelled = false;
loop {
// Race between getting the next event and cancellation
let event = futures::select! {
// Race between getting the first event and cancellation
let first_event = futures::select! {
event = events.next().fuse() => event,
_ = cancellation_rx.changed().fuse() => {
if *cancellation_rx.borrow() {
@@ -1482,25 +1482,54 @@ impl Thread {
continue;
}
};
let Some(event) = event else {
let Some(first_event) = first_event else {
break;
};
log::trace!("Received completion event: {:?}", event);
match event {
Ok(event) => {
tool_results.extend(this.update(cx, |this, cx| {
this.handle_completion_event(
event,
event_stream,
cancellation_rx.clone(),
cx,
)
})??);
}
Err(err) => {
error = Some(err);
break;
// Collect all immediately available events to process as a batch
let mut batch = vec![first_event];
while let Some(event) = events.next().now_or_never().flatten() {
batch.push(event);
}
// Process the batch in a single update
let batch_result = this.update(cx, |this, cx| {
let mut batch_tool_results = Vec::new();
let mut batch_error = None;
for event in batch {
log::trace!("Received completion event: {:?}", event);
match event {
Ok(event) => {
match this.handle_completion_event(
event,
event_stream,
cancellation_rx.clone(),
cx,
) {
Ok(Some(task)) => batch_tool_results.push(task),
Ok(None) => {}
Err(err) => {
batch_error = Some(err);
break;
}
}
}
Err(err) => {
batch_error = Some(err.into());
break;
}
}
}
cx.notify();
(batch_tool_results, batch_error)
})?;
tool_results.extend(batch_result.0);
if let Some(err) = batch_result.1 {
error = Some(err.downcast()?);
break;
}
}
@@ -1638,11 +1667,11 @@ impl Thread {
self.flush_pending_message(cx);
self.pending_message = Some(AgentMessage::default());
}
Text(new_text) => self.handle_text_event(new_text, event_stream, cx),
Text(new_text) => self.handle_text_event(new_text, event_stream),
Thinking { text, signature } => {
self.handle_thinking_event(text, signature, event_stream, cx)
self.handle_thinking_event(text, signature, event_stream)
}
RedactedThinking { data } => self.handle_redacted_thinking_event(data, cx),
RedactedThinking { data } => self.handle_redacted_thinking_event(data),
ReasoningDetails(details) => {
let last_message = self.pending_message();
// Store the last non-empty reasoning_details (overwrites earlier ones)
@@ -1702,12 +1731,7 @@ impl Thread {
Ok(None)
}
fn handle_text_event(
&mut self,
new_text: String,
event_stream: &ThreadEventStream,
cx: &mut Context<Self>,
) {
fn handle_text_event(&mut self, new_text: String, event_stream: &ThreadEventStream) {
event_stream.send_text(&new_text);
let last_message = self.pending_message();
@@ -1718,8 +1742,6 @@ impl Thread {
.content
.push(AgentMessageContent::Text(new_text));
}
cx.notify();
}
fn handle_thinking_event(
@@ -1727,7 +1749,6 @@ impl Thread {
new_text: String,
new_signature: Option<String>,
event_stream: &ThreadEventStream,
cx: &mut Context<Self>,
) {
event_stream.send_thinking(&new_text);
@@ -1743,16 +1764,13 @@ impl Thread {
signature: new_signature,
});
}
cx.notify();
}
fn handle_redacted_thinking_event(&mut self, data: String, cx: &mut Context<Self>) {
fn handle_redacted_thinking_event(&mut self, data: String) {
let last_message = self.pending_message();
last_message
.content
.push(AgentMessageContent::RedactedThinking(data));
cx.notify();
}
fn handle_tool_use_event(