use crate::{
door::route::{event, sent},
serve::storage::Failure,
};
use serde_json::{Value, json};
use std::{collections::BTreeMap, io::Write};
#[derive(Default)]
struct Durable {
through: u64,
ranges: Vec<(u64, u64)>,
}
impl Durable {
fn contains(&self, seq: u64) -> bool {
seq <= self.through
|| self
.ranges
.iter()
.any(|&(first, last)| first <= seq && seq <= last)
}
fn insert(&mut self, seq: u64) {
let at = self.ranges.partition_point(|&(_, last)| last < seq);
self.ranges.insert(at, (seq, seq));
let mut i = 0;
while i + 1 < self.ranges.len() {
if self.ranges[i].1.saturating_add(1) >= self.ranges[i + 1].0 {
let next = self.ranges.remove(i + 1);
self.ranges[i].1 = self.ranges[i].1.max(next.1);
} else {
i += 1;
}
}
if self
.ranges
.first()
.is_some_and(|&(first, _)| first == self.through.saturating_add(1))
{
self.through = self.ranges.remove(0).1;
}
}
}
pub(super) struct Output<'a> {
writer: &'a mut dyn Write,
conversation: Option<&'a str>,
cursors: BTreeMap<String, Durable>,
bound: usize,
}
impl<'a> Output<'a> {
pub fn new(
writer: &'a mut dyn Write,
conversation: Option<&'a str>,
after: u64,
bound: usize,
) -> Self {
let mut cursors = BTreeMap::new();
if let Some(id) = conversation {
cursors.insert(
id.to_owned(),
Durable {
through: after,
ranges: Vec::new(),
},
);
}
Self {
writer,
conversation,
cursors,
bound,
}
}
pub fn through(&self, conversation: &str) -> u64 {
self.cursors
.get(conversation)
.map_or(0, |cursor| cursor.through)
}
pub fn item(&mut self, item: &Value) -> Result<(), Failure> {
let mut prefix = String::new();
if let (Some(conversation), Some(seq)) =
(item["conversation"].as_str(), item["seq"].as_u64())
{
let other = self
.conversation
.is_some_and(|selected| selected != conversation);
if other
&& !matches!(
item["kind"].as_str(),
Some("session_opening" | "session_opened")
)
{
return Ok(());
}
if !other {
let cursor = self.cursors.entry(conversation.to_owned()).or_default();
if cursor.contains(seq) {
return Ok(());
}
if item["durable"] == true {
let previous = cursor.through;
cursor.insert(seq);
if cursor.through > previous {
prefix = format!("id: {conversation}:{}\n", cursor.through);
}
}
let bytes = self.cursors.iter().fold(0usize, |bytes, (id, c)| {
bytes
.saturating_add(id.len())
.saturating_add(std::mem::size_of::<Durable>())
.saturating_add(
c.ranges
.len()
.saturating_mul(std::mem::size_of::<(u64, u64)>()),
)
});
if bytes > self.bound {
return Err(crate::door::assistant_state::bound(
"events_after_bytes",
self.bound,
bytes,
));
}
}
}
prefix.push_str(&event("item", &item.to_string()));
self.send(&prefix)
}
pub fn replayed(&mut self, conversation: &str) -> Result<(), Failure> {
self.send(&event(
"replayed",
&json!({"conversation":conversation,"through":self.through(conversation)}).to_string(),
))
}
fn send(&mut self, frame: &str) -> Result<(), Failure> {
match sent(self.writer, frame) {
Ok(true) => Ok(()),
Ok(false) => Err(Failure::new(
"assistant_stream_closed",
"the browser closed its stream",
)),
Err(error) => Err(Failure::new(
"assistant_stream_write_failed",
error.to_string(),
)),
}
}
}