Skip to content

Commit 87a70ad

Browse files
fix: accept upstream SSE data without a space (#269)
## Summary An upstream SSE field may omit the space after `data:`, but the live HTTP reader currently discards that spelling before event normalization. A valid unspaced stream consequently produces an empty completed Responses output over HTTP and WebSocket, and the Messages loop misses its tool call. An unspaced `data:[DONE]` also fails to stop a connection that remains open. Accept both spellings in the shared inference reader and retain the existing `data: ` form for its consumers. Preserve the field value, including additional whitespace, and recognize both exact completion markers. This completes the live-transport boundary missing from the decoder-only coverage in #236; the [SSE field grammar makes the space optional](https://html.spec.whatwg.org/multipage/server-sent-events.html#parsing-an-event-stream). Production changes are confined to `response_lines`: no new parser, lifecycle state, public type, database schema, dependency, or delivery queue. Already accepted data lines remain unchanged. Add coverage through real HTTP/WebSocket clients, SQLite close/reopen and a subsequent upstream continuation request, plus recorded Messages tool rounds. Correct the search test fixture to use GET and verify the successful tool result reaches the next model request. ## Test Plan Exact commit `74b28f187a57f8d0bd41c2f884ea919b21a378f4` verified from a clean detached worktree with a separate build directory, Rust 1.98.0, and macOS: - Against unchanged production on main `60126bb`, the final tests yield seven intended failures and eleven passing controls. The isolated patched build passes all eighteen. HTTP and WebSocket failures show `response.completed` with `output: []`; Messages stops after one round instead of two; the reader drops unspaced fields or times out after the unspaced marker. - Cover mixed spacing, empty values, Unicode, LF/CRLF, case-sensitive field names, malformed JSON passthrough, meaningful extra whitespace, exact and near-miss completion markers, the 256 KiB line boundary, and chunk timeout behavior. Existing split-UTF-8, error, cancellation, slow-consumer, and lifecycle tests also pass. - `cargo test --locked --workspace --all-features`: 1,086 passed; eight PostgreSQL tests and one existing doctest ignored. - Forty-five fresh-process repeat runs pass: 40 reader boundary tests, 70 Messages tests, and 75 HTTP/WebSocket restart tests, including three concurrent test processes. - Workspace formatting, all-target/all-feature check and Clippy with warnings denied; binary builds, launcher contracts, and source-install CLI E2E with a development-profile wheel pass. - Python: 111 passed, three installation-only skips. Cassette validation: 113/113, 258 turns. No recorded cassette files were modified. - All applicable changed-file hooks pass. Remaining repository-wide hooks pass; the known Apple Git hang in the all-files large-file hook was avoided by running that hook on the complete changed-file set. Validation uses local HTTP/database fixtures and existing recorded streams. No live-model/GPU or new PostgreSQL execution is claimed. This is an optional-space fix, not a claim of complete SSE framing conformance. Signed-off-by: Chuyue Wang <stevenwang0805@outlook.com> Co-authored-by: Francisco Javier Arceo <arceofrancisco@gmail.com>
1 parent 04ce90e commit 87a70ad

3 files changed

Lines changed: 381 additions & 10 deletions

File tree

‎crates/agentic-server-core/src/executor/inference.rs‎

Lines changed: 131 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -190,7 +190,8 @@ pub fn call_inference(
190190
}
191191
}
192192

193-
/// Convert a successful upstream response body into normalized SSE data lines.
193+
/// Convert a successful upstream response body into SSE data lines with a `data: ` prefix.
194+
/// The single space after the field colon is optional on the wire.
194195
pub(super) fn response_lines(
195196
resp: reqwest::Response,
196197
chunk_timeout: Duration,
@@ -215,12 +216,19 @@ pub(super) fn response_lines(
215216
return;
216217
}
217218
};
218-
for line in lines {
219-
match line.as_str() {
220-
"data: [DONE]" => return,
221-
l if l.starts_with("data: ") => yield Ok(line),
222-
_ => {}
219+
for mut line in lines {
220+
let Some(data) = line.strip_prefix("data:") else {
221+
continue;
222+
};
223+
if data.strip_prefix(' ').unwrap_or(data) == "[DONE]" {
224+
return;
225+
}
226+
// Keep one transport spelling for Responses and Messages consumers,
227+
// without trimming whitespace that belongs to the field value.
228+
if !data.starts_with(' ') {
229+
line.insert("data:".len(), ' ');
223230
}
231+
yield Ok(line);
224232
}
225233
}
226234
}
@@ -239,6 +247,123 @@ mod tests {
239247

240248
use super::*;
241249

250+
async fn read_sse_body(body: String, keep_open: bool) -> Vec<ExecutorResult<String>> {
251+
let app = axum::Router::new().route(
252+
"/v1/responses",
253+
post(move || {
254+
let body = body.clone();
255+
async move {
256+
let chunks = async_stream::stream! {
257+
yield Ok::<_, Infallible>(Bytes::from(body));
258+
if keep_open {
259+
std::future::pending::<()>().await;
260+
}
261+
};
262+
Response::builder()
263+
.header("content-type", "text/event-stream")
264+
.body(Body::from_stream(chunks))
265+
.unwrap()
266+
}
267+
}),
268+
);
269+
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
270+
let url = format!("http://{}/v1/responses", listener.local_addr().unwrap());
271+
let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
272+
let result = call_inference(
273+
"{}".to_owned(),
274+
url,
275+
Arc::new(reqwest::Client::new()),
276+
None,
277+
Duration::from_secs(1),
278+
)
279+
.collect()
280+
.await;
281+
server.abort();
282+
let _ = server.await;
283+
result
284+
}
285+
286+
#[tokio::test]
287+
async fn sse_data_spacing_preserves_field_values() {
288+
let lines = read_sse_body(
289+
concat!(
290+
":comment\r\nevent: ignored\r\nid: 1\r\nretry: 1000\r\n",
291+
"Data: ignored\r\ndatabase: ignored\r\n data: ignored\r\n",
292+
"data:{\"delta\":\"雪 ☃: data: text\"}\r\n\r\n",
293+
"data: {\"spaced\":true}\n\n",
294+
"data: {\"leading_space\":true}\n\n",
295+
"data:\t{\"tab\":true}\n\n",
296+
"data:\n\ndata: \n\ndata:{malformed}\n\n",
297+
"data: [DONE]\n\ndata:[DONE]extra\n\n",
298+
)
299+
.to_owned(),
300+
false,
301+
)
302+
.await
303+
.into_iter()
304+
.collect::<ExecutorResult<Vec<_>>>()
305+
.unwrap();
306+
assert_eq!(
307+
lines,
308+
[
309+
"data: {\"delta\":\"雪 ☃: data: text\"}",
310+
"data: {\"spaced\":true}",
311+
"data: {\"leading_space\":true}",
312+
"data: \t{\"tab\":true}",
313+
"data: ",
314+
"data: ",
315+
"data: {malformed}",
316+
"data: [DONE]",
317+
"data: [DONE]extra",
318+
]
319+
);
320+
}
321+
322+
#[tokio::test]
323+
async fn sse_data_spacing_stops_at_both_done_markers() {
324+
for separator in ["", " "] {
325+
let body = format!("data: {{}}\n\ndata:{separator}[DONE]\r\n\r\ndata: ignored\n\n");
326+
let lines = read_sse_body(body, true)
327+
.await
328+
.into_iter()
329+
.collect::<ExecutorResult<Vec<_>>>()
330+
.expect("[DONE] must stop without waiting for upstream EOF or a chunk timeout");
331+
assert_eq!(lines, ["data: {}"]);
332+
}
333+
}
334+
335+
#[tokio::test]
336+
async fn sse_data_spacing_preserves_line_limit() {
337+
for separator in ["", " "] {
338+
let prefix = format!("data:{separator}");
339+
let line = format!("{prefix}{}", "x".repeat(MAX_SSE_LINE_BYTES - prefix.len()));
340+
let accepted = read_sse_body(format!("{line}\n\n"), false).await;
341+
assert_eq!(accepted.len(), 1);
342+
assert_eq!(
343+
accepted[0].as_ref().unwrap().len(),
344+
MAX_SSE_LINE_BYTES + usize::from(separator.is_empty())
345+
);
346+
347+
let rejected = read_sse_body(format!("{line}x\n\n"), false).await;
348+
assert_eq!(rejected.len(), 1);
349+
assert!(
350+
rejected[0]
351+
.as_ref()
352+
.unwrap_err()
353+
.to_string()
354+
.contains("upstream SSE line exceeded")
355+
);
356+
}
357+
}
358+
359+
#[tokio::test]
360+
async fn sse_data_spacing_preserves_chunk_timeout_without_done() {
361+
let lines = read_sse_body("data:{}\n\n".to_owned(), true).await;
362+
assert_eq!(lines.len(), 2);
363+
assert_eq!(lines[0].as_ref().unwrap(), "data: {}");
364+
assert!(lines[1].as_ref().unwrap_err().to_string().contains("chunk timeout"));
365+
}
366+
242367
async fn oversized_body_server(status: StatusCode) -> (String, tokio::task::JoinHandle<()>) {
243368
let app = axum::Router::new().route(
244369
"/v1/responses",

‎crates/agentic-server-core/tests/messages_stream_test.rs‎

Lines changed: 39 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ use agentic_core::tool::{ToolRegistry, WebSearchHandler};
1919
use agentic_core::types::messages::{GatewayToolMap, ToolParam, registry_tools};
2020
use axum::extract::State;
2121
use axum::response::{IntoResponse, Response};
22-
use axum::routing::post;
22+
use axum::routing::{get, post};
2323
use axum::{Json, Router};
2424
use futures::StreamExt;
2525
use http::StatusCode;
@@ -146,7 +146,7 @@ async fn spawn_mock_vllm_stream_then_error(
146146
async fn spawn_mock_search() -> (String, tokio::task::JoinHandle<()>) {
147147
let app = Router::new().route(
148148
"/v1/search",
149-
post(|Json(_body): Json<Value>| async move {
149+
get(|| async move {
150150
Json(serde_json::json!({
151151
"results": {"web": [{"url": "https://www.rust-lang.org/", "title": "Rust",
152152
"description": "d", "snippets": ["Rust 1.89.0 is the latest stable release."]}], "news": []},
@@ -188,8 +188,28 @@ async fn run_test_messages_stream(
188188

189189
#[tokio::test]
190190
async fn messages_stream_presents_one_message_and_hides_gateway_tool() {
191-
let (vllm_url, upstream, _v) = spawn_mock_vllm_stream(cassette_turn_streams()).await;
192-
let (search_url, _s) = spawn_mock_search().await;
191+
assert_messages_stream_presents_one_message(cassette_turn_streams()).await;
192+
}
193+
194+
#[tokio::test]
195+
async fn messages_stream_accepts_unspaced_sse_data_through_gateway_tool_rounds() {
196+
let streams = cassette_turn_streams()
197+
.into_iter()
198+
.map(|body| {
199+
body.split_inclusive('\n')
200+
.map(|line| {
201+
line.strip_prefix("data: ")
202+
.map_or_else(|| line.to_owned(), |data| format!("data:{data}"))
203+
})
204+
.collect()
205+
})
206+
.collect();
207+
assert_messages_stream_presents_one_message(streams).await;
208+
}
209+
210+
async fn assert_messages_stream_presents_one_message(streams: Vec<String>) {
211+
let (vllm_url, upstream, vllm) = spawn_mock_vllm_stream(streams).await;
212+
let (search_url, search) = spawn_mock_search().await;
193213
let exec_ctx = build_exec_ctx(&vllm_url, &search_url).await;
194214

195215
let request = serde_json::json!({
@@ -210,13 +230,28 @@ async fn messages_stream_presents_one_message_and_hides_gateway_tool() {
210230
let stream = run_test_messages_stream(request, registry, Arc::clone(&exec_ctx)).await;
211231
let chunks: Vec<String> = stream.collect().await;
212232
let sse = chunks.join("");
233+
vllm.abort();
234+
search.abort();
235+
let _ = tokio::join!(vllm, search);
213236

214237
// Two upstream rounds ran (tool round + final).
215238
assert_eq!(
216239
upstream.calls.load(Ordering::SeqCst),
217240
2,
218241
"one tool round + one final round"
219242
);
243+
let requests = upstream.requests.lock().await;
244+
let tool_output = &requests[1]["messages"].as_array().unwrap().last().unwrap()["content"][0];
245+
assert_eq!(tool_output["type"], "tool_result");
246+
assert_ne!(
247+
tool_output["is_error"], true,
248+
"gateway search must succeed: {tool_output}"
249+
);
250+
assert!(
251+
tool_output["content"]
252+
.to_string()
253+
.contains("https://www.rust-lang.org/")
254+
);
220255

221256
// Exactly one logical message lifecycle.
222257
assert_eq!(

0 commit comments

Comments
 (0)