Skip to content

Commit 84bd606

Browse files
aaronvgcursoragent
aaronvg
and
cursoragent
authored
Prioritize latest sse events for llm streaming (#2467)
# Pull Request Template Thanks for taking the time to fill out this pull request! ## Issue Reference Please link to any related issues - [ ] This PR fixes/closes #[issue number] ## Changes Please describe the changes proposed in this pull request This PR modifies the `orchestrate_stream` function in `engine/baml-runtime/src/internal/llm_client/orchestrator/stream.rs` to prioritize parsing the latest received Server-Sent Event (SSE) content from LLM streams. Instead of sequentially parsing every buffered SSE event, the updated logic: - Decouples SSE ingestion from partial parsing. - Maintains only the latest `LLMResponse::Success` snapshot and its content string. - Uses a `tokio::select` loop with a zero-duration sleep to opportunistically parse and emit the *most recent* content when there's a lull in incoming events, effectively skipping older, buffered partial responses. This change aims to reduce perceived latency by ensuring that the UI or downstream consumers always receive the freshest available partial parse, rather than processing potentially stale intermediate states. ## Testing Please describe how you tested these changes - [ ] Unit tests added/updated - [x] Manual testing performed (via `cargo test --lib` in `engine` to ensure compilation and existing tests pass) - [ ] Tested in [environment] ## Screenshots If applicable, add screenshots to help explain your changes [Add screenshots here...] ## PR Checklist Please ensure you've completed these items - [x] I have read and followed the contributing guidelines - [x] My code follows the style guidelines of this project - [x] I have performed a self-review of my own code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings ## Additional Notes Add any other context about the PR here The primary goal of this change is to reduce latency in LLM streaming by ensuring that only the most recent SSE event's content is parsed and emitted. This prevents the system from getting bogged down processing intermediate, quickly outdated partial responses, leading to a more responsive user experience. --- <a href="https://cursor.com/background-agent?bcId=bc-629d4483-4c3d-49b4-bdc8-1854eed935a7"> <picture> <source media="(prefers-color-scheme: dark)" srcset="https://cursor.com/open-in-cursor-dark.svg"> <source media="(prefers-color-scheme: light)" srcset="https://cursor.com/open-in-cursor-light.svg"> <img alt="Open in Cursor" src="https://cursor.com/open-in-cursor.svg"> </picture> </a> <a href="https://cursor.com/agents?id=bc-629d4483-4c3d-49b4-bdc8-1854eed935a7"> <picture> <source media="(prefers-color-scheme: dark)" srcset="https://cursor.com/open-in-web-dark.svg"> <source media="(prefers-color-scheme: light)" srcset="https://cursor.com/open-in-web-light.svg"> <img alt="Open in Web" src="https://cursor.com/open-in-web.svg"> </picture> </a> --------- Co-authored-by: Cursor Agent <cursoragent@cursor.com>
1 parent bedb929 commit 84bd606

8 files changed

Lines changed: 309 additions & 79 deletions

File tree

‎engine/baml-runtime/src/internal/llm_client/orchestrator/stream.rs‎

Lines changed: 123 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,13 @@ use std::sync::Arc;
33
use anyhow::Result;
44
use async_std::stream::StreamExt;
55
use baml_ids::HttpRequestId;
6-
use baml_types::{BamlValue, BamlValueWithMeta};
6+
use baml_types::BamlValue;
7+
use futures::StreamExt as FuturesStreamExt;
78
use internal_baml_core::ir::repr::IntermediateRepr;
89
use jsonish::BamlValueWithFlags;
910
use serde_json::json;
1011
use stream_cancel::Tripwire;
12+
use tokio::time::MissedTickBehavior;
1113
use web_time::Duration;
1214

1315
use super::{call::CtxWithHttpRequestId, OrchestrationScope, OrchestratorNodeIterator};
@@ -96,32 +98,128 @@ where
9698
let ctx = CtxWithHttpRequestId::from(ctx);
9799
let stream_res = node.stream(&ctx, &prompt).await;
98100
let final_response = match stream_res {
99-
Ok(response) => response
100-
.map(|stream_part| {
101-
if let Some(on_tick) = on_tick_fn.as_ref() {
102-
on_tick();
101+
Ok(mut response_stream) => {
102+
let mut last_response: Option<LLMResponse> = None;
103+
let mut latest_success_snapshot: Option<crate::internal::llm_client::LLMCompleteResponse> = None;
104+
let mut latest_content_for_parse: Option<String> = None;
105+
// Track last parsed payload surfaced to downstream listeners so we can dedupe events
106+
let mut last_sent_partial_serialized: Option<String> = None;
107+
let mut parse_interval = tokio::time::interval(std::time::Duration::from_millis(20));
108+
// If parsing falls behind, skip missed ticks so we only parse latest.
109+
parse_interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
110+
111+
loop {
112+
tokio::select! {
113+
// Prioritize consuming SSE events over parsing.
114+
biased;
115+
maybe_item = FuturesStreamExt::next(&mut response_stream) => {
116+
match maybe_item {
117+
Some(stream_part) => {
118+
if let Some(on_tick) = on_tick_fn.as_ref() {
119+
on_tick();
120+
}
121+
match &stream_part {
122+
LLMResponse::Success(s) => {
123+
// Track latest snapshot and content
124+
latest_success_snapshot = Some(s.clone());
125+
latest_content_for_parse = Some(s.content.clone());
126+
last_response = Some(LLMResponse::Success(s.clone()));
127+
}
128+
other => {
129+
last_response = Some(other.clone());
130+
}
131+
}
132+
}
133+
None => {
134+
// End of stream
135+
break;
136+
}
137+
}
138+
}
139+
// Periodically surface the latest partial parse to downstream listeners.
140+
_ = parse_interval.tick(), if on_event.is_some() => {
141+
if let Some(on_event) = on_event.as_ref() {
142+
if let Some(snap) = latest_success_snapshot.as_ref() {
143+
if let Some(mut content) = latest_content_for_parse.take() {
144+
match partial_parse_fn(&content) {
145+
Ok(baml_value) => {
146+
// Strip flags to reduce memory usage
147+
let parsed = ResponseBamlValue(baml_value.0.map_meta_owned(|m| {
148+
jsonish::ResponseValueMeta(vec![], m.1, m.2, m.3)
149+
}));
150+
if let Ok(serialized) = serde_json::to_string(&parsed.serialize_partial()) {
151+
if last_sent_partial_serialized
152+
.as_deref()
153+
!= Some(serialized.as_str())
154+
{
155+
// only successful events sent to the client
156+
on_event(FunctionResult::new(
157+
node.scope.clone(),
158+
LLMResponse::Success(snap.clone()),
159+
Some(Ok(parsed)),
160+
));
161+
last_sent_partial_serialized = Some(serialized);
162+
}
163+
} else {
164+
// If serialization fails, still emit the parsed event instead of dropping it.
165+
on_event(FunctionResult::new(
166+
node.scope.clone(),
167+
LLMResponse::Success(snap.clone()),
168+
Some(Ok(parsed)),
169+
));
170+
// Intentionally do not update last_sent_partial_serialized here.
171+
}
172+
}
173+
Err(_) => {
174+
// Only restore the content if nothing newer has arrived since we took it.
175+
if latest_content_for_parse.is_none() {
176+
latest_content_for_parse = Some(content);
177+
}
178+
}
179+
}
180+
}
181+
}
182+
}
183+
}
103184
}
104-
if let Some(on_event) = on_event.as_ref() {
105-
if let LLMResponse::Success(s) = &stream_part {
106-
let response_value = partial_parse_fn(&s.content);
107-
// Flags seem to use a ton of memory, so we strip them here.
108-
if let Ok(baml_value) = response_value {
109-
// only success events are sent to the stream
110-
on_event(FunctionResult::new(
111-
node.scope.clone(),
112-
LLMResponse::Success(s.clone()),
113-
Some(Ok(ResponseBamlValue(baml_value.0.map_meta_owned(|m| {
114-
jsonish::ResponseValueMeta(vec![], m.1, m.2, m.3)
115-
})))),
116-
));
185+
}
186+
187+
if let Some(on_event) = on_event.as_ref() {
188+
if let Some(snap) = latest_success_snapshot.as_ref() {
189+
if let Some(mut content) = latest_content_for_parse.take() {
190+
if let Ok(baml_value) = partial_parse_fn(&content) {
191+
// Strip flags to reduce memory usage
192+
let parsed = ResponseBamlValue(baml_value.0.map_meta_owned(|m| {
193+
jsonish::ResponseValueMeta(vec![], m.1, m.2, m.3)
194+
}));
195+
if let Ok(serialized) = serde_json::to_string(&parsed.serialize_partial()) {
196+
if last_sent_partial_serialized
197+
.as_deref()
198+
!= Some(serialized.as_str())
199+
{
200+
// Only successful events should reach downstream listeners
201+
on_event(FunctionResult::new(
202+
node.scope.clone(),
203+
LLMResponse::Success(snap.clone()),
204+
Some(Ok(parsed)),
205+
));
206+
last_sent_partial_serialized = Some(serialized);
207+
}
208+
} else {
209+
// If serialization fails, still emit the parsed event instead of dropping it.
210+
on_event(FunctionResult::new(
211+
node.scope.clone(),
212+
LLMResponse::Success(snap.clone()),
213+
Some(Ok(parsed)),
214+
));
215+
// Intentionally do not update last_sent_partial_serialized here.
216+
}
117217
}
118218
}
119219
}
120-
stream_part
121-
})
122-
.fold(None, |_, current| Some(current))
123-
.await
124-
.unwrap_or_else(|| {
220+
}
221+
222+
last_response.unwrap_or_else(|| {
125223
LLMResponse::LLMFailure(LLMErrorResponse {
126224
client: node.provider.name().into(),
127225
model: None,
@@ -132,7 +230,8 @@ where
132230
message: "Stream ended without response".to_string(),
133231
code: crate::internal::llm_client::ErrorCode::from_u16(2),
134232
})
135-
}),
233+
})
234+
}
136235
Err(response) => response,
137236
};
138237

‎engine/baml-runtime/src/types/response.rs‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,8 +139,16 @@ impl FunctionResult {
139139

140140
// Capture the actual error to preserve its details
141141
let actual_error = err.to_string();
142+
142143
// TODO: HACK! Figure out why now connection errors dont get converted into ExposedError. Instead of converting to a validation error, check for connection errors here. We probably are missing a lot of other connection failures that should NOT be validation errors.
143-
if actual_error.to_lowercase().contains("connecterror") {
144+
if actual_error.to_lowercase().contains("connecterror")
145+
|| actual_error
146+
.to_lowercase()
147+
.contains("profilefile provider could not be built")
148+
|| actual_error
149+
.to_lowercase()
150+
.contains("session token not found")
151+
{
144152
return ExposedError::ClientHttpError {
145153
client_name: match self.llm_response() {
146154
LLMResponse::Success(resp) => resp.client.clone(),

‎engine/generators/languages/python/src/_templates/runtime.py.j2‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -216,6 +216,8 @@ class {{ call_manager_name }}:
216216
# on_tick
217217
# always None! sync streams don't support on_tick
218218
None,
219+
# abort_controller
220+
resolved_options.abort_controller,
219221
)
220222
return ctx, result
221223

‎engine/language_client_python/python_src/baml_py/stream.py‎

Lines changed: 45 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -81,12 +81,32 @@ async def __aiter__(self):
8181
while True:
8282
try:
8383
event = self.__event_queue.get_nowait()
84-
if event is None:
85-
break
86-
if event.is_ok():
87-
yield self.__partial_coerce(event)
8884
except queue.Empty:
89-
await asyncio.sleep(0.050)
85+
await asyncio.sleep(0.010)
86+
continue
87+
88+
if event is None:
89+
break
90+
91+
# Drain the queue to coalesce and keep only the most recent successful event.
92+
latest_ok: Optional[FunctionResult] = event if event.is_ok() else None
93+
done_seen = False
94+
while True:
95+
try:
96+
nxt = self.__event_queue.get_nowait()
97+
if nxt is None:
98+
done_seen = True
99+
break
100+
if nxt.is_ok():
101+
latest_ok = nxt
102+
except queue.Empty:
103+
break
104+
105+
if latest_ok is not None:
106+
yield self.__partial_coerce(latest_ok)
107+
108+
if done_seen:
109+
break
90110
except Exception as e:
91111
raise e
92112
finally:
@@ -163,8 +183,26 @@ def __iter__(self):
163183
event = self.__event_queue.get()
164184
if event is None:
165185
break
166-
if event.is_ok():
167-
yield self.__partial_coerce(event)
186+
187+
# Drain the queue to coalesce and keep only the most recent successful event.
188+
latest_ok: Optional[FunctionResult] = event if event.is_ok() else None
189+
done_seen = False
190+
while True:
191+
try:
192+
nxt = self.__event_queue.get_nowait()
193+
if nxt is None:
194+
done_seen = True
195+
break
196+
if nxt.is_ok():
197+
latest_ok = nxt
198+
except queue.Empty:
199+
break
200+
201+
if latest_ok is not None:
202+
yield self.__partial_coerce(latest_ok)
203+
204+
if done_seen:
205+
break
168206
except Exception as e:
169207
raise e
170208
finally:

‎integ-tests/python-v1/baml_client/runtime.py‎

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎integ-tests/python/baml_client/runtime.py‎

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)