use std::collections::VecDeque;
use std::sync::{Arc, Mutex};
use serde::Serialize;
use tokio::sync::broadcast;
pub(crate) type SseBroadcast = broadcast::Sender<Arc<String>>;
pub(crate) const SSE_HEARTBEAT_INTERVAL: std::time::Duration = std::time::Duration::from_secs(15);
pub(crate) const SSE_HEARTBEAT_COMMENT: &[u8] = b":heartbeat\n\n";
#[derive(Debug, Clone, Serialize)]
pub(crate) struct TaskStatusUpdateEvent {
pub id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub context_id: Option<String>,
pub status: StreamStatus,
pub r#final: bool,
}
#[derive(Debug, Clone, Serialize)]
pub(crate) struct StreamStatus {
pub state: String,
pub message: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub response: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub reopen: Option<u32>,
}
#[derive(Debug, Clone, Serialize)]
pub(crate) struct TaskArtifactUpdateEvent {
pub id: String,
pub artifact: StreamArtifact,
}
#[derive(Debug, Clone, Serialize)]
pub(crate) struct StreamArtifact {
pub tool_name: String,
pub content: String,
pub fence_source: Option<String>,
pub tool_call_id: Option<String>,
pub is_error: bool,
pub duration_ms: Option<u64>,
pub exit_code: Option<i32>,
pub truncated: bool,
}
impl StreamArtifact {
#[allow(clippy::too_many_arguments)]
pub(crate) fn tool_call(
tool_name: String,
content: String,
fence_source: Option<String>,
tool_call_id: &str,
is_error: bool,
duration_ms: u64,
exit_code: Option<i32>,
truncated: bool,
) -> Self {
Self {
tool_name,
content,
fence_source,
tool_call_id: (!tool_call_id.is_empty()).then(|| tool_call_id.to_string()),
is_error,
duration_ms: Some(duration_ms),
exit_code,
truncated,
}
}
pub(crate) fn hook(hook_name: String, payload: String) -> Self {
Self {
tool_name: hook_name,
content: payload,
fence_source: None,
tool_call_id: None,
is_error: false,
duration_ms: None,
exit_code: None,
truncated: false,
}
}
}
pub(crate) fn format_sse_event(event_id: u64, event_type: &str, data: &str) -> String {
format!("id: {event_id}\nevent: {event_type}\ndata: {data}\n\n")
}
pub(crate) fn format_unnumbered_sse_event(event_type: &str, data: &str) -> String {
format!("event: {event_type}\ndata: {data}\n\n")
}
pub(crate) fn frame_id(frame: &str) -> Option<u64> {
frame.strip_prefix("id: ")?.split('\n').next()?.parse().ok()
}
pub(crate) fn admit_live_frame(last_written: &mut Option<u64>, frame: &str) -> bool {
let Some(id) = frame_id(frame) else {
return true;
};
if last_written.is_some_and(|written| id <= written) {
return false;
}
*last_written = Some(id);
true
}
pub(crate) enum ReplayResult {
Complete(Vec<Arc<String>>),
WithGap {
first_available_id: u64,
events: Vec<Arc<String>>,
},
}
pub(crate) fn format_gap_event(first_available_id: u64) -> String {
format_unnumbered_sse_event(
"gap",
&format!("{{\"first_available_id\":{first_available_id}}}"),
)
}
pub(crate) fn format_lagged_event(missed: u64) -> String {
format!("event: lagged\ndata: {{\"missed\":{missed}}}\n\n")
}
pub(crate) struct SseEventBuffer {
events: VecDeque<(u64, Arc<String>)>,
capacity: usize,
next_id: u64,
}
impl SseEventBuffer {
pub fn new(capacity: usize) -> Self {
Self {
events: VecDeque::new(),
capacity,
next_id: 1,
}
}
fn append(&mut self, event_type: &str, data: &str) -> (u64, Arc<String>) {
let id = self.next_id;
self.next_id += 1;
let event = Arc::new(format_sse_event(id, event_type, data));
if self.events.len() >= self.capacity {
self.events.pop_front();
}
self.events.push_back((id, Arc::clone(&event)));
(id, event)
}
pub fn replay_from(&self, last_id: u64) -> ReplayResult {
let Some(oldest_id) = self.events.front().map(|(id, _)| *id) else {
return ReplayResult::Complete(vec![]);
};
if last_id >= self.next_id || last_id < oldest_id - 1 {
let events = self.events.iter().map(|(_, s)| Arc::clone(s)).collect();
return ReplayResult::WithGap {
first_available_id: oldest_id,
events,
};
}
let events = self
.events
.iter()
.filter(|(id, _)| *id > last_id)
.map(|(_, s)| Arc::clone(s))
.collect();
ReplayResult::Complete(events)
}
}
pub(crate) fn emit_frame(
tx: &SseBroadcast,
buf: &Arc<Mutex<SseEventBuffer>>,
event_type: &str,
data: &str,
) -> u64 {
let mut buf = buf.lock().unwrap();
let (id, event) = buf.append(event_type, data);
let _ = tx.send(event);
id
}
pub(crate) async fn emit_sse(
sse: &Option<(SseBroadcast, Arc<Mutex<SseEventBuffer>>)>,
event_type: &str,
data: &impl Serialize,
) {
let Some((sse_tx, sse_buffer)) = sse else {
return;
};
let data_str = match serde_json::to_string(data) {
Ok(s) => s,
Err(_) => return,
};
emit_frame(sse_tx, sse_buffer, event_type, &data_str);
}
pub(crate) fn is_final_status_for(event: &str, task_id: &str) -> bool {
let mut is_status = false;
let mut data = None;
for line in event.lines() {
if line == "event: status" {
is_status = true;
} else if let Some(rest) = line.strip_prefix("data: ") {
data = Some(rest);
}
}
if !is_status {
return false;
}
let Some(parsed) = data.and_then(|data| serde_json::from_str::<serde_json::Value>(data).ok())
else {
return false;
};
parsed.get("final") == Some(&serde_json::Value::Bool(true))
&& parsed.get("id").and_then(serde_json::Value::as_str) == Some(task_id)
}
pub(crate) fn emit_chunk_sse(
tx: &SseBroadcast,
buf: &Arc<Mutex<SseEventBuffer>>,
task_id: &str,
chunk: &str,
) {
emit_frame(
tx,
buf,
"text",
&format_text_sse_data(task_id, chunk, false),
);
}
pub(crate) fn emit_thinking_chunk_sse(
tx: &SseBroadcast,
buf: &Arc<Mutex<SseEventBuffer>>,
task_id: &str,
chunk: &str,
) {
emit_frame(
tx,
buf,
"thinking",
&format_text_sse_data(task_id, chunk, false),
);
}
pub(crate) fn emit_chunk_sse_final(
tx: &SseBroadcast,
buf: &Arc<Mutex<SseEventBuffer>>,
task_id: &str,
text: &str,
) {
emit_frame(tx, buf, "text", &format_text_sse_data(task_id, text, true));
}
fn format_text_sse_data(task_id: &str, text: &str, is_final: bool) -> String {
let text_json = serde_json::to_string(text).unwrap_or_else(|_| "\"\"".to_string());
format!(r#"{{"id":"{task_id}","text":{text_json},"final":{is_final}}}"#)
}
#[cfg(test)]
mod tests {
use super::*;
fn frame_json(artifact: StreamArtifact) -> String {
serde_json::to_string(&TaskArtifactUpdateEvent {
id: "task_1".into(),
artifact,
})
.unwrap()
}
#[test]
fn tool_call_frame_serializes_every_key_in_declaration_order() {
let json = frame_json(StreamArtifact::tool_call(
"bash".into(),
"hello".into(),
Some("tool:bash".into()),
"call_1",
false,
12,
Some(0),
false,
));
assert_eq!(
json,
r#"{"id":"task_1","artifact":{"tool_name":"bash","content":"hello","fence_source":"tool:bash","tool_call_id":"call_1","is_error":false,"duration_ms":12,"exit_code":0,"truncated":false}}"#
);
}
#[test]
fn tool_call_frame_passes_truncated_through() {
let json = frame_json(StreamArtifact::tool_call(
"read".into(),
"partial".into(),
None,
"call_1",
false,
3,
None,
true,
));
assert!(json.contains(r#""truncated":true"#), "{json}");
}
#[test]
fn tool_call_frame_keeps_the_exit_code_key_with_or_without_a_subprocess() {
let ran = frame_json(StreamArtifact::tool_call(
"bash".into(),
String::new(),
None,
"call_1",
true,
1204,
Some(2),
false,
));
assert!(ran.contains(r#""exit_code":2"#), "{ran}");
assert!(ran.contains(r#""is_error":true"#), "{ran}");
let none = frame_json(StreamArtifact::tool_call(
"write-file".into(),
String::new(),
None,
"call_1",
true,
3,
None,
false,
));
assert!(none.contains(r#""exit_code":null"#), "{none}");
}
#[test]
fn tool_call_frame_records_an_empty_provider_id_as_null() {
let empty = frame_json(StreamArtifact::tool_call(
"bash".into(),
String::new(),
None,
"",
false,
1,
None,
false,
));
assert!(empty.contains(r#""tool_call_id":null"#), "{empty}");
let issued = frame_json(StreamArtifact::tool_call(
"bash".into(),
String::new(),
None,
"toolu_01",
false,
1,
None,
false,
));
assert!(issued.contains(r#""tool_call_id":"toolu_01""#), "{issued}");
}
#[test]
fn hook_frame_reports_no_call_behind_it() {
let json = frame_json(StreamArtifact::hook(
"my-hook".into(),
r#"{"reviewed":true}"#.into(),
));
assert_eq!(
json,
r#"{"id":"task_1","artifact":{"tool_name":"my-hook","content":"{\"reviewed\":true}","fence_source":null,"tool_call_id":null,"is_error":false,"duration_ms":null,"exit_code":null,"truncated":false}}"#
);
}
#[test]
fn format_sse_event_produces_correct_frame() {
let out = format_sse_event(3, "status", r#"{"id":"x"}"#);
assert_eq!(out, "id: 3\nevent: status\ndata: {\"id\":\"x\"}\n\n");
}
fn buffer_with(capacity: usize, count: u64) -> SseEventBuffer {
let mut buf = SseEventBuffer::new(capacity);
for n in 1..=count {
buf.append("status", &n.to_string());
}
buf
}
fn ids(events: &[Arc<String>]) -> Vec<u64> {
events.iter().map(|e| frame_id(e).unwrap()).collect()
}
fn complete(result: ReplayResult) -> Vec<u64> {
match result {
ReplayResult::Complete(events) => ids(&events),
ReplayResult::WithGap { .. } => panic!("expected Complete"),
}
}
fn with_gap(result: ReplayResult) -> (u64, Vec<u64>) {
match result {
ReplayResult::WithGap {
first_available_id,
events,
} => (first_available_id, ids(&events)),
ReplayResult::Complete(_) => panic!("expected WithGap"),
}
}
#[test]
fn ids_start_at_one_and_rise_by_one() {
let (tx, mut rx) = broadcast::channel(8);
let buf = Arc::new(Mutex::new(SseEventBuffer::new(8)));
assert_eq!(emit_frame(&tx, &buf, "status", "{}"), 1);
emit_chunk_sse(&tx, &buf, "t", "a");
emit_thinking_chunk_sse(&tx, &buf, "t", "b");
emit_chunk_sse_final(&tx, &buf, "t", "");
assert_eq!(emit_frame(&tx, &buf, "artifact", "{}"), 5);
let received: Vec<u64> = std::iter::from_fn(|| rx.try_recv().ok())
.map(|e| frame_id(&e).unwrap())
.collect();
assert_eq!(received, vec![1, 2, 3, 4, 5]);
}
#[test]
fn replay_after_eviction_writes_a_gap() {
let buf = buffer_with(3, 5);
assert_eq!(with_gap(buf.replay_from(1)), (3, vec![3, 4, 5]));
assert_eq!(with_gap(buf.replay_from(0)), (3, vec![3, 4, 5]));
assert_eq!(complete(buf.replay_from(2)), vec![3, 4, 5]);
}
#[test]
fn replay_from_zero_without_eviction_is_every_frame() {
let buf = buffer_with(8, 5);
assert_eq!(complete(buf.replay_from(0)), vec![1, 2, 3, 4, 5]);
assert_eq!(complete(buf.replay_from(3)), vec![4, 5]);
assert_eq!(complete(buf.replay_from(5)), Vec::<u64>::new());
}
#[test]
fn replay_from_a_cursor_never_issued_writes_a_gap() {
assert_eq!(
with_gap(buffer_with(8, 5).replay_from(99)),
(1, vec![1, 2, 3, 4, 5])
);
assert_eq!(
with_gap(buffer_with(3, 5).replay_from(6)),
(3, vec![3, 4, 5])
);
}
#[test]
fn replay_from_an_empty_buffer_is_empty() {
let buf = SseEventBuffer::new(8);
assert_eq!(complete(buf.replay_from(0)), Vec::<u64>::new());
assert_eq!(complete(buf.replay_from(99)), Vec::<u64>::new());
}
#[test]
fn concurrent_emitters_share_one_ascending_sequence() {
const THREADS: u64 = 8;
const PER_THREAD: u64 = 200;
let total = THREADS * PER_THREAD;
let (tx, mut rx) = broadcast::channel(total as usize);
let buf = Arc::new(Mutex::new(SseEventBuffer::new(total as usize)));
std::thread::scope(|scope| {
for t in 0..THREADS {
let (tx, buf) = (tx.clone(), Arc::clone(&buf));
scope.spawn(move || {
for _ in 0..PER_THREAD {
if t % 2 == 0 {
emit_chunk_sse(&tx, &buf, "t", "x");
} else {
emit_frame(&tx, &buf, "status", "{}");
}
}
});
}
});
let buffered = complete(buf.lock().unwrap().replay_from(0));
let received: Vec<u64> = std::iter::from_fn(|| rx.try_recv().ok())
.map(|e| frame_id(&e).unwrap())
.collect();
let expected: Vec<u64> = (1..=total).collect();
assert_eq!(buffered, expected, "buffer order is id order");
assert_eq!(received, expected, "live delivery order is id order");
}
#[test]
fn frame_id_reads_the_leading_id_line() {
assert_eq!(frame_id("id: 42\nevent: status\ndata: {}\n\n"), Some(42));
assert_eq!(frame_id("event: gap\ndata: {}\n\n"), None);
assert_eq!(frame_id(":heartbeat\n\n"), None);
}
#[test]
fn live_frames_already_replayed_are_dropped() {
let frame = |id| format_sse_event(id, "status", "{}");
let mut last = Some(3);
assert!(!admit_live_frame(&mut last, &frame(2)));
assert!(!admit_live_frame(&mut last, &frame(3)));
assert!(admit_live_frame(&mut last, &frame(4)));
assert_eq!(last, Some(4));
assert!(
!admit_live_frame(&mut last, &frame(4)),
"an id is written once"
);
assert!(admit_live_frame(&mut last, &format_gap_event(1)));
assert_eq!(last, Some(4), "a frame without an id leaves the cursor");
let mut fresh = None;
assert!(admit_live_frame(&mut fresh, &frame(1)));
assert_eq!(fresh, Some(1));
}
#[test]
fn unnumbered_frame_has_no_id_line() {
let out = format_unnumbered_sse_event("status", r#"{"id":"x"}"#);
assert_eq!(out, "event: status\ndata: {\"id\":\"x\"}\n\n");
assert_eq!(frame_id(&out), None);
}
#[test]
fn format_gap_event_has_correct_shape() {
let s = format_gap_event(42);
assert!(s.contains("event: gap\n"), "missing event field");
assert!(s.contains("\"first_available_id\":42"), "missing id field");
assert!(s.ends_with("\n\n"), "missing double newline terminator");
}
#[test]
fn format_lagged_event_has_exact_wire_form() {
let s = format_lagged_event(7);
assert_eq!(s, "event: lagged\ndata: {\"missed\":7}\n\n");
let data = s
.lines()
.find_map(|line| line.strip_prefix("data: "))
.expect("data line");
let parsed: serde_json::Value = serde_json::from_str(data).expect("data is JSON");
assert_eq!(parsed, serde_json::json!({"missed": 7}));
}
#[test]
fn is_final_status_for_matches_only_this_tasks_final_status() {
let status_final = "id: 5\nevent: status\ndata: {\"id\":\"x\",\"final\":true}\n\n";
let other_final = "id: 6\nevent: status\ndata: {\"id\":\"xy\",\"final\":true}\n\n";
let text_final =
"id: 3\nevent: text\ndata: {\"id\":\"x\",\"text\":\"\",\"final\":true}\n\n";
let status_nonfinal = "id: 1\nevent: status\ndata: {\"id\":\"x\",\"final\":false}\n\n";
assert!(
is_final_status_for(status_final, "x"),
"this task's final status is terminal"
);
assert!(
!is_final_status_for(other_final, "x"),
"another task's final status is not this task's end, even when its id contains this one"
);
assert!(
!is_final_status_for(status_final, "xy"),
"an id that only contains the frame's id does not match it"
);
assert!(
!is_final_status_for(text_final, "x"),
"text:final:true is not terminal"
);
assert!(
!is_final_status_for(status_nonfinal, "x"),
"status:final:false is not terminal"
);
}
}