diff --git a/rust/src/bench/src/backends/streaming.rs b/rust/src/bench/src/backends/streaming.rs index 16837a274a..90ec05cbc4 100644 --- a/rust/src/bench/src/backends/streaming.rs +++ b/rust/src/bench/src/backends/streaming.rs @@ -6,7 +6,7 @@ /// Accumulates incoming byte chunks and extracts complete SSE messages. /// Mirrors Python's `StreamedResponseHandler` from endpoint_request_func.py:22-60. pub struct StreamedResponseHandler { - buffer: String, + buffer: Vec, /// Reusable message buffer — avoids allocating a new Vec per `add_chunk` call. messages: Vec, } @@ -14,7 +14,7 @@ pub struct StreamedResponseHandler { impl StreamedResponseHandler { pub fn new() -> Self { Self { - buffer: String::with_capacity(4096), + buffer: Vec::with_capacity(4096), messages: Vec::with_capacity(4), } } @@ -26,12 +26,11 @@ impl StreamedResponseHandler { pub fn add_chunk(&mut self, chunk_bytes: &[u8]) -> &[String] { self.messages.clear(); - let chunk_str = String::from_utf8_lossy(chunk_bytes); - self.buffer.push_str(&chunk_str); + self.buffer.extend_from_slice(chunk_bytes); // Split by double newlines (SSE message separator) - while let Some(pos) = self.buffer.find("\n\n") { - let message = self.buffer[..pos].trim().to_string(); + while let Some(pos) = self.buffer.windows(2).position(|window| window == b"\n\n") { + let message = String::from_utf8_lossy(&self.buffer[..pos]).trim().to_string(); // Efficiently remove consumed bytes by shifting remaining data self.buffer.drain(..pos + 2); if !message.is_empty() { @@ -47,19 +46,20 @@ impl StreamedResponseHandler { // // Also handles multi-field SSE events where the buffer may start with // "event: ...\ndata: ..." (Dynamo frontend). - let data_start = if self.buffer.starts_with("data: ") { + let buffer = String::from_utf8_lossy(&self.buffer); + let data_start = if buffer.starts_with("data: ") { Some(0) } else { // Look for a "data: " line in multi-field events - self.buffer.find("\ndata: ").map(|p| p + 1) + buffer.find("\ndata: ").map(|p| p + 1) }; if let Some(offset) = data_start { - let content = self.buffer[offset + 6..].trim(); + let content = buffer[offset + 6..].trim(); if content == "[DONE]" || (!content.is_empty() && serde_json::from_str::<&serde_json::value::RawValue>(content).is_ok()) { - self.messages.push(self.buffer.trim().to_string()); + self.messages.push(buffer.trim().to_string()); self.buffer.clear(); } } @@ -102,6 +102,26 @@ mod tests { assert_eq!(msgs2.len(), 1); } + #[test] + fn test_split_multibyte_utf8() { + let mut handler = StreamedResponseHandler::new(); + let message = "data: {\"test\":\"中😀\"}\n\n"; + let mut messages = Vec::new(); + + for chunk in message.as_bytes().chunks(1) { + messages.extend_from_slice(handler.add_chunk(chunk)); + } + + assert_eq!(messages, &[message.trim().to_string()]); + } + + #[test] + fn test_invalid_utf8_remains_lossy() { + let mut handler = StreamedResponseHandler::new(); + let msgs = handler.add_chunk(b"data: {\"test\":\"\xFF\"}\n\n"); + assert_eq!(msgs, &["data: {\"test\":\"�\"}".to_string()]); + } + #[test] fn test_comment_lines() { let mut handler = StreamedResponseHandler::new();