fix(backend): improve sse parsing (#7100)

backend/Cargo.lock

@@ -5144,6 +5144,17 @@ dependencies = [
  "pin-project-lite",
 ]
 
+[[package]]
+name = "eventsource-stream"
+version = "0.2.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "74fef4569247a5f429d9156b9d0a2599914385dd189c539334c625d8099d90ab"
+dependencies = [
+ "futures-core",
+ "nom 7.1.3",
+ "pin-project-lite",
+]
+
 [[package]]
 name = "fallible-iterator"
 version = "0.2.0"
@@ -15777,6 +15788,7 @@ dependencies = [
  "deno_webidl",
  "dotenv",
  "dyn-iter",
+ "eventsource-stream",
  "flume",
  "futures",
  "gcp_auth",

backend/Cargo.toml

@@ -252,6 +252,7 @@ urlencoding = "^2"
 url = { version = "^2" , features = ["serde"]}
 async-oauth2 = "0.5.1"
 reqwest = { version = "^0.12", features = ["json", "stream", "gzip", "multipart"] }
+eventsource-stream = "0.2.3"
 time = "^0"
 serde_urlencoded = "^0"
 tokio-tar = "^0"

backend/windmill-worker/Cargo.toml

@@ -115,6 +115,7 @@ nix.workspace = true
 bytes.workspace = true
 reqwest.workspace = true
 reqwest-middleware.workspace = true
+eventsource-stream.workspace = true
 mime_guess.workspace = true
 hex.workspace = true
 tiberius = { workspace = true, optional = true }

backend/windmill-worker/src/ai/sse.rs

@@ -1,5 +1,6 @@
 use std::collections::HashMap;
 
+use eventsource_stream::Eventsource;
 use reqwest::Response;
 use serde::Deserialize;
 use serde_json;
@@ -51,38 +52,23 @@ pub trait SSEParser {
     async fn parse_event_data(&mut self, data: &str) -> Result<(), Error>;
 
     async fn parse_events(&mut self, response: Response) -> Result<(), Error> {
-        let mut stream = response.bytes_stream();
-        let mut buffer = String::new();
+        let mut stream = response.bytes_stream().eventsource();
 
-        while let Some(chunk_result) = stream.next().await {
-            let chunk = chunk_result
-                .map_err(|e| Error::internal_err(format!("Failed to read chunk: {}", e)))?;
+        while let Some(event) = stream.next().await {
+            match event {
+                Ok(event) => {
+                    if *DEBUG_SSE_STREAM {
+                        tracing::info!("SSE event: {:?}", event);
+                    }
 
-            // Convert chunk to string and add to buffer
-            let chunk_str = String::from_utf8_lossy(&chunk);
-            if *DEBUG_SSE_STREAM {
-                tracing::info!("SSE chunk: {}", chunk_str);
-            }
-            buffer.push_str(&chunk_str);
-
-            // Process complete lines from buffer
-            while let Some(newline_pos) = buffer.find("\n\n") {
-                let line = buffer.drain(..newline_pos + 2).collect::<String>();
-                let line = line.trim_end_matches('\n');
-
-                // Skip empty lines and comments
-                if line.is_empty() || line.starts_with(':') {
-                    continue;
-                }
-
-                // Parse SSE data field
-                if let Some(data) = line.strip_prefix("data: ") {
-                    if data == "[DONE]" {
-                        // OpenAI sends [DONE] to indicate end of stream
+                    if event.data == "[DONE]" {
                         return Ok(());
                     }
 
-                    self.parse_event_data(data).await?;
+                    self.parse_event_data(&event.data).await?;
+                }
+                Err(e) => {
+                    tracing::error!("Failed to parse SSE event: {}", e);
                 }
             }
         }