Files
ai-app/app/src/client/event_stream.rs
T

126 lines
4.1 KiB
Rust

use std::io::{BufRead, BufReader};
use event_model::SeqEvent;
use crate::client::api::{ApiError, Transport};
use crate::client::sse::SseReader;
/// The frame name the server uses to say a cursor was too far behind to
/// continue from. Must match `send_backlog` in `server/src/routes.rs`.
const RESET_EVENT: &str = "reset";
/// One frame of a session's event stream, folded from the wire shape the
/// caller needs to act on -- mirroring what `EventStream.kt`'s three
/// callbacks were for, as a single enum instead, since Rust has no
/// equivalent of handing three closures to one blocking call.
pub enum StreamItem {
Open,
Reset,
Event { raw: String, event: SeqEvent },
}
/// Follows `/sessions/{id}/events?after={after}`, calling `on_item` for
/// each [`StreamItem`] until the connection drops or `on_item` asks to
/// stop (by returning `false`). Reconnecting -- with the last seq seen as
/// the new cursor -- is the caller's job.
pub fn follow_session_events(
transport: &dyn Transport,
session_id: &str,
after: u64,
mut on_item: impl FnMut(StreamItem) -> bool,
) -> Result<(), ApiError> {
let path = format!("/sessions/{session_id}/events?after={after}");
let body = transport.stream(&path)?;
if !on_item(StreamItem::Open) {
return Ok(());
}
let mut lines = BufReader::new(body).lines();
let mut reader = SseReader::new();
while let Some(line) = lines.next().transpose().map_err(|e| ApiError {
message: format!("Can't reach the server -- retrying. ({e})"),
status: None,
})? {
let Some(frame) = reader.feed_line(&line) else {
continue;
};
if frame.name.as_deref() == Some(RESET_EVENT) {
if !on_item(StreamItem::Reset) {
return Ok(());
}
} else if !frame.data.is_empty() {
let event: SeqEvent = serde_json::from_str(&frame.data).map_err(|e| ApiError {
message: format!("The server sent an event this build couldn't parse: {e}"),
status: None,
})?;
if !on_item(StreamItem::Event {
raw: frame.data,
event,
}) {
return Ok(());
}
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::client::api::{Body, RawResponse};
use std::io::Cursor;
struct FixtureTransport {
body: &'static str,
}
impl Transport for FixtureTransport {
fn request(
&self,
_method: &str,
_path: &str,
_body: Option<Body>,
) -> Result<RawResponse, ApiError> {
unimplemented!("this fixture only serves a stream")
}
fn stream(&self, _path: &str) -> Result<Box<dyn std::io::Read + Send>, ApiError> {
Ok(Box::new(Cursor::new(self.body.as_bytes().to_vec())))
}
}
#[test]
fn events_and_a_reset_frame_are_told_apart() {
let transport = FixtureTransport {
body: "event:reset\n\ndata:{\"seq\":1,\"ts\":1.0,\"type\":\"status\",\"state\":\"idle\"}\n\n",
};
let mut items = Vec::new();
follow_session_events(&transport, "s1", 0, |item| {
items.push(match item {
StreamItem::Open => "open".to_string(),
StreamItem::Reset => "reset".to_string(),
StreamItem::Event { event, .. } => format!("event:{}", event.seq),
});
true
})
.unwrap();
assert_eq!(items, vec!["open", "reset", "event:1"]);
}
#[test]
fn the_caller_can_stop_early() {
let transport = FixtureTransport {
body: "data:{\"seq\":1,\"ts\":1.0,\"type\":\"status\",\"state\":\"idle\"}\n\n\
data:{\"seq\":2,\"ts\":1.0,\"type\":\"status\",\"state\":\"idle\"}\n\n",
};
let mut count = 0;
follow_session_events(&transport, "s1", 0, |item| {
if matches!(item, StreamItem::Event { .. }) {
count += 1;
}
count < 1
})
.unwrap();
assert_eq!(count, 1);
}
}