[WIP] SSE partially working

Author: wdoppenberg

rust-sdk/Cargo.lock

@@ -14,7 +14,9 @@ futures = "0.3.31"
 json-patch = "4.0.0"
 log = "0.4.27"
 reqwest = { version = "0.12.22" , features = ["json", "stream"]}
-sse-client = "1.1.1"
+bytes = "1.5.0"
+tokio = "1.36.0"
 
 [dev-dependencies]
+env_logger = "0.11.8"
 tokio = { version = "1.36.0", features = ["full"] }

rust-sdk/crates/ag-ui-client/Cargo.toml

@@ -0,0 +1,40 @@
+use ag_ui_client::sse::SseResponseExt;
+use ag_ui_core::event::Event;
+use futures::StreamExt;
+use serde::Deserialize;
+use std::error::Error;
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "lowercase")]
+enum EventType {
+    Ping,
+    Update,
+    Message,
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn Error>> {
+    // Create a client
+    let client = reqwest::Client::new();
+
+    // Example 1: Stream with custom event and data types
+    println!("Example 1: Typed events with custom event and data types");
+    let response = client.get("https://httpbun.org/sse").send().await?;
+    let mut stream = response.event_source().await;
+
+    while let Some(result) = stream.next().await {
+        match result {
+            Ok(sse_event) => {
+                if let Some(event_type) = &sse_event.event {
+                    match event_type.as_str() {
+                        "ping" => println!("Ping: {}", sse_event.data),
+                        &_ => panic!("Unknown event type {event_type}")
+                    }
+                }
+            },
+            Err(err) => eprintln!("Error: {}", err),
+        }
+    }
+    Ok(())
+}
+

rust-sdk/crates/ag-ui-client/examples/sse_example.rs

@@ -1,15 +1,15 @@
-use std::string::ParseError;
 use crate::Agent;
 use crate::agent::AgentError;
+use crate::sse::SseResponseExt;
 use crate::stream::EventStream;
-use ag_ui_core::event::Event;
+use ag_ui_core::event::{Event, EventType};
 use ag_ui_core::types::input::RunAgentInput;
 use ag_ui_core::{AgentState, FwdProps};
 use async_trait::async_trait;
 use futures::StreamExt;
+use log::{debug, trace};
 use reqwest::header::HeaderMap;
 use reqwest::{Client as HttpClient, Url};
-use sse_client::EventSource;
 
 pub struct HttpAgent {
     http_client: HttpClient,
@@ -42,30 +42,31 @@ impl<StateT: AgentState, FwdPropsT: FwdProps> Agent<StateT, FwdPropsT> for HttpA
         &self,
         input: &RunAgentInput<StateT, FwdPropsT>,
     ) -> Result<EventStream<'async_trait, StateT>, AgentError> {
-        let event_source = EventSource::new(&self.base_url.as_str())
-            .map_err(|e| AgentError::ExecutionError {
-                message: e.to_string(),
-            })?;
-
-        let stream = self
+        // Send the request and get the response
+        let response = self
             .http_client
             .post(self.base_url.clone())
             .json(input)
             .headers(self.header_map.clone())
             .send()
-            .await?
-            .bytes_stream()
-            .map(|result| {
-                result.map_err(Into::into).and_then(|bytes| {
-                    serde_json::from_slice::<Event<StateT>>(&bytes).map_err(|e| {
-                        AgentError::ExecutionError {
-                            message: e.to_string(),
-                        }
-                    })
-                })
+            .await?;
+
+        // Convert the response to an SSE event stream
+        let stream = response
+            .event_source()
+            .await
+            .map(|result| match result {
+                Ok(event) => {
+                    trace!("Received event: {event:?}");
+                    let event_data: Event<StateT> = serde_json::from_str(&event.data)?;
+                    trace!("Deserialized event: {event_data:?}");
+                    Ok(event_data)
+                }
+                Err(err) => Err(AgentError::ExecutionError {
+                    message: err.to_string(),
+                }),
             })
             .boxed();
-
         Ok(stream)
     }
 }

rust-sdk/crates/ag-ui-client/src/http.rs

@@ -3,6 +3,8 @@ pub mod event_handler;
 pub mod http;
 mod stream;
 mod subscriber;
+pub mod sse;
 
 pub use agent::Agent;
 pub use http::HttpAgent;
+pub use sse::SseResponseExt;