use a2a_protocol_types::{JsonRpcResponse, StreamResponse};
use hyper::body::Bytes;
use tokio::sync::mpsc;
use tokio::task::AbortHandle;
use crate::error::{ClientError, ClientResult};
use crate::streaming::sse_parser::SseParser;
pub(crate) type BodyChunk = ClientResult<Bytes>;
const EVENT_BRIDGE_CAPACITY: usize = 64;
pub struct EventStream {
rx: mpsc::Receiver<BodyChunk>,
parser: SseParser,
phase: Phase,
abort_handle: Option<AbortHandle>,
status_code: u16,
jsonrpc_envelope: bool,
first_event_timeout: Option<std::time::Duration>,
idle_timeout: Option<std::time::Duration>,
last_activity: tokio::time::Instant,
first_chunk_received: bool,
at_final_event: bool,
events_received: u64,
last_event_id: Option<String>,
held: Option<
Box<dyn std::any::Any + Send + Sync + std::panic::UnwindSafe + std::panic::RefUnwindSafe>,
>,
}
impl EventStream {
#[must_use]
#[cfg(any(test, feature = "websocket"))]
pub(crate) fn new(rx: mpsc::Receiver<BodyChunk>) -> Self {
Self {
rx,
parser: SseParser::new(),
phase: Phase::Open,
abort_handle: None,
status_code: 200,
jsonrpc_envelope: true,
first_event_timeout: None,
idle_timeout: Some(crate::config::DEFAULT_STREAM_IDLE_TIMEOUT),
last_activity: tokio::time::Instant::now(),
first_chunk_received: false,
at_final_event: false,
events_received: 0,
last_event_id: None,
held: None,
}
}
#[must_use]
#[cfg(test)]
pub(crate) fn with_abort_handle(
rx: mpsc::Receiver<BodyChunk>,
abort_handle: AbortHandle,
) -> Self {
Self {
rx,
parser: SseParser::new(),
phase: Phase::Open,
abort_handle: Some(abort_handle),
status_code: 200,
jsonrpc_envelope: true,
first_event_timeout: None,
idle_timeout: Some(crate::config::DEFAULT_STREAM_IDLE_TIMEOUT),
last_activity: tokio::time::Instant::now(),
first_chunk_received: false,
at_final_event: false,
events_received: 0,
last_event_id: None,
held: None,
}
}
#[must_use]
pub fn from_event_channel(mut rx: mpsc::Receiver<ClientResult<StreamResponse>>) -> Self {
let (tx, body_rx) = mpsc::channel::<BodyChunk>(EVENT_BRIDGE_CAPACITY);
let bridge = tokio::spawn(async move {
while let Some(event) = rx.recv().await {
let chunk = match event {
Ok(ref ev) => serde_json::to_string(ev).map_or_else(
|e| Err(ClientError::Serialization(e)),
|json| {
Ok(Bytes::from(format!(
"data: {{\"jsonrpc\":\"2.0\",\"id\":null,\"result\":{json}}}\n\n"
)))
},
),
Err(e) => Err(e),
};
if tx.send(chunk).await.is_err() {
break;
}
}
});
Self::with_status(body_rx, bridge.abort_handle(), 200)
}
#[must_use]
pub(crate) fn with_status(
rx: mpsc::Receiver<BodyChunk>,
abort_handle: AbortHandle,
status_code: u16,
) -> Self {
Self {
rx,
parser: SseParser::new(),
phase: Phase::Open,
abort_handle: Some(abort_handle),
status_code,
jsonrpc_envelope: true,
first_event_timeout: None,
idle_timeout: Some(crate::config::DEFAULT_STREAM_IDLE_TIMEOUT),
last_activity: tokio::time::Instant::now(),
first_chunk_received: false,
at_final_event: false,
events_received: 0,
last_event_id: None,
held: None,
}
}
#[must_use]
pub(crate) const fn with_jsonrpc_envelope(mut self, envelope: bool) -> Self {
self.jsonrpc_envelope = envelope;
self
}
#[must_use]
pub const fn with_first_event_timeout(mut self, timeout: std::time::Duration) -> Self {
self.first_event_timeout = Some(timeout);
self
}
#[must_use]
pub const fn with_idle_timeout(mut self, timeout: Option<std::time::Duration>) -> Self {
self.idle_timeout = timeout;
self
}
#[must_use]
pub const fn with_max_event_size(mut self, max_bytes: usize) -> Self {
self.parser.set_max_event_size(max_bytes);
self
}
#[must_use]
#[cfg(feature = "websocket")]
pub(crate) fn holding(
mut self,
resource: impl std::any::Any
+ Send
+ Sync
+ std::panic::UnwindSafe
+ std::panic::RefUnwindSafe
+ 'static,
) -> Self {
self.held = Some(Box::new(resource));
self
}
#[must_use]
pub const fn status_code(&self) -> u16 {
self.status_code
}
#[must_use]
pub fn last_event_id(&self) -> Option<&str> {
self.last_event_id.as_deref()
}
pub async fn next(&mut self) -> Option<ClientResult<StreamResponse>> {
loop {
if let Some(result) = self.parser.next_frame() {
match result {
Ok(frame) => {
self.last_event_id = frame.id;
return Some(self.decode_frame(&frame.data));
}
Err(e) => {
return Some(Err(ClientError::Transport(e.to_string())));
}
}
}
match self.phase {
Phase::Finished => return None,
Phase::BodyEnded => {
self.phase = Phase::Finished;
return self.judge_end().map(Err);
}
Phase::Open => {}
}
let chunk = match self.silence_deadline() {
Some((deadline, bound)) => {
let Ok(chunk) = tokio::time::timeout_at(deadline, self.rx.recv()).await else {
return Some(Err(self.fail_silent(bound)));
};
chunk
}
None => self.rx.recv().await,
};
match chunk {
None => {
self.phase = Phase::BodyEnded;
}
Some(Err(e)) => {
self.phase = Phase::Finished;
return Some(Err(e));
}
Some(Ok(bytes)) => {
self.first_chunk_received = true;
self.last_activity = tokio::time::Instant::now();
self.parser.feed(&bytes);
}
}
}
}
fn silence_deadline(&self) -> Option<(tokio::time::Instant, std::time::Duration)> {
let bound = if self.first_chunk_received {
self.idle_timeout
} else {
self.first_event_timeout
}?;
Some((self.last_activity.checked_add(bound)?, bound))
}
fn fail_silent(&mut self, bound: std::time::Duration) -> ClientError {
self.phase = Phase::Finished;
if let Some(handle) = self.abort_handle.take() {
handle.abort();
}
if self.first_chunk_received {
ClientError::Timeout(format!(
"stream idle timeout: no data, not even a keep-alive comment, for {bound:?} \
since the last data received; the task on the server is not cancelled, \
so resubscribe to continue"
))
} else {
ClientError::Timeout(format!(
"stream produced no data before the first-event timeout ({bound:?})"
))
}
}
fn judge_end(&self) -> Option<ClientError> {
let partial = self.parser.incomplete_event_len();
if self.at_final_event {
if let Some(bytes) = partial {
trace_warn!(
bytes,
"SSE body ended mid-frame after the stream's final event; \
discarding the unterminated frame"
);
let _ = bytes;
}
return None;
}
let detail = match (partial, self.events_received) {
(Some(bytes), _) => format!(
"the connection closed mid-frame, discarding {bytes} bytes of an \
unterminated event"
),
(None, 0) => "the body ended before any event".to_owned(),
(None, n) => format!(
"the body ended after {n} event(s), none of them a final one \
(a Message, or a terminal or interrupted task state)"
),
};
Some(ClientError::IncompleteStream {
last_event_id: self.last_event_id.clone(),
detail,
})
}
const fn record(&mut self, event: &StreamResponse) {
self.events_received += 1;
self.at_final_event = ends_stream(event);
}
fn decode_frame(&mut self, data: &str) -> ClientResult<StreamResponse> {
if self.jsonrpc_envelope {
let envelope: JsonRpcResponse<StreamResponse> =
serde_json::from_str(data).map_err(ClientError::Serialization)?;
match envelope {
JsonRpcResponse::Success(ok) => {
self.record(&ok.result);
if is_terminal(&ok.result) {
self.phase = Phase::Finished;
}
Ok(ok.result)
}
JsonRpcResponse::Error(err) => {
self.phase = Phase::Finished;
let a2a = crate::transport::map_jsonrpc_error(
err.error.code,
err.error.message,
err.error.data,
);
Err(ClientError::Protocol(a2a))
}
}
} else {
let event: StreamResponse = match serde_json::from_str(data) {
Ok(event) => event,
Err(e) => {
return Err(
match crate::transport::rest::error_frame::decode_stream_error_frame(data) {
Some(a2a) => {
self.phase = Phase::Finished;
ClientError::Protocol(a2a)
}
None => ClientError::Serialization(e),
},
);
}
};
self.record(&event);
if is_terminal(&event) {
self.phase = Phase::Finished;
}
Ok(event)
}
}
}
impl Drop for EventStream {
fn drop(&mut self) {
if let Some(handle) = self.abort_handle.take() {
handle.abort();
}
drop(self.held.take());
}
}
#[allow(clippy::missing_fields_in_debug)]
impl std::fmt::Debug for EventStream {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("EventStream")
.field("done", &(self.phase == Phase::Finished))
.field("pending_frames", &self.parser.pending_count())
.finish()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Phase {
Open,
BodyEnded,
Finished,
}
const fn ends_stream(event: &StreamResponse) -> bool {
let state = match event {
StreamResponse::Message(_) => return true,
StreamResponse::Task(task) => task.status.state,
StreamResponse::StatusUpdate(update) => update.status.state,
_ => return false,
};
state.is_terminal() || state.is_interrupted()
}
const fn is_terminal(event: &StreamResponse) -> bool {
matches!(
event,
StreamResponse::StatusUpdate(ev) if ev.status.state.is_terminal()
)
}
#[cfg(test)]
const fn _event_stream_is_unwind_safe() {
const fn assert_unwind_safe<T: std::panic::UnwindSafe + std::panic::RefUnwindSafe>() {}
assert_unwind_safe::<EventStream>();
}
#[cfg(test)]
mod tests {
use super::*;
use a2a_protocol_types::{
JsonRpcSuccessResponse, JsonRpcVersion, TaskId, TaskState, TaskStatus,
TaskStatusUpdateEvent,
};
use std::time::Duration;
const TEST_TIMEOUT: Duration = Duration::from_secs(5);
fn make_status_event(state: TaskState, _is_final: bool) -> StreamResponse {
StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
task_id: TaskId::new("t1"),
context_id: a2a_protocol_types::ContextId::new("c1"),
status: TaskStatus {
state,
message: None,
timestamp: None,
},
metadata: None,
})
}
fn sse_frame(event: &StreamResponse) -> String {
let resp = JsonRpcSuccessResponse {
jsonrpc: JsonRpcVersion,
id: Some(serde_json::json!(1)),
result: event.clone(),
};
let json = serde_json::to_string(&resp).unwrap();
format!("data: {json}\n\n")
}
#[tokio::test]
async fn stream_delivers_events() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx);
let event = make_status_event(TaskState::Working, false);
let sse_bytes = sse_frame(&event);
tx.send(Ok(Bytes::from(sse_bytes))).await.unwrap();
drop(tx);
let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out")
.unwrap()
.unwrap();
assert!(
matches!(result, StreamResponse::StatusUpdate(ref ev) if ev.status.state == TaskState::Working)
);
}
#[tokio::test]
async fn stream_ends_on_final_event() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx);
let event = make_status_event(TaskState::Completed, true);
let sse_bytes = sse_frame(&event);
tx.send(Ok(Bytes::from(sse_bytes))).await.unwrap();
let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out waiting for final event")
.unwrap()
.unwrap();
assert!(
matches!(result, StreamResponse::StatusUpdate(ref ev) if ev.status.state == TaskState::Completed)
);
let end = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out waiting for stream end");
assert!(end.is_none());
}
#[tokio::test]
async fn stream_propagates_body_error() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx);
tx.send(Err(ClientError::Transport("network error".into())))
.await
.unwrap();
let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out")
.unwrap();
assert!(result.is_err());
}
#[tokio::test]
async fn a_channel_closed_before_any_event_is_incomplete() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx);
drop(tx);
let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out");
match result {
Some(Err(ClientError::IncompleteStream {
last_event_id,
detail,
})) => {
assert_eq!(last_event_id, None);
assert_eq!(detail, "the body ended before any event");
}
other => panic!("expected IncompleteStream, got {other:?}"),
}
assert!(
stream.next().await.is_none(),
"reported once, then finished"
);
}
#[tokio::test]
async fn the_end_of_the_body_is_judged_by_the_last_event() {
for (states, clean) in [
(vec![TaskState::Working, TaskState::InputRequired], true),
(vec![TaskState::Working, TaskState::AuthRequired], true),
(vec![TaskState::InputRequired, TaskState::Working], false),
(vec![TaskState::Submitted, TaskState::Working], false),
] {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx).with_jsonrpc_envelope(false);
for state in &states {
let ev = make_status_event(*state, false);
tx.send(Ok(Bytes::from(bare_sse_frame(&ev)))).await.unwrap();
}
drop(tx);
for _ in &states {
assert!(matches!(stream.next().await, Some(Ok(_))), "{states:?}");
}
let end = stream.next().await;
if clean {
assert!(end.is_none(), "{states:?}: a completion, got {end:?}");
} else {
match end {
Some(Err(ClientError::IncompleteStream { detail, .. })) => {
assert!(detail.contains("after 2 event(s)"), "{detail}");
}
other => panic!("{states:?}: expected IncompleteStream, got {other:?}"),
}
}
assert!(stream.next().await.is_none());
}
}
#[test]
fn ends_stream_covers_final_events_only() {
use a2a_protocol_types::{Message, MessageId, MessageRole, Part, Task};
for state in [
TaskState::Completed,
TaskState::Failed,
TaskState::Canceled,
TaskState::Rejected,
TaskState::InputRequired,
TaskState::AuthRequired,
] {
assert!(ends_stream(&make_status_event(state, false)), "{state:?}");
}
for state in [TaskState::Submitted, TaskState::Working] {
assert!(!ends_stream(&make_status_event(state, false)), "{state:?}");
}
let task = |state| {
StreamResponse::Task(Task {
id: TaskId::new("t"),
context_id: a2a_protocol_types::ContextId::new("c"),
status: TaskStatus::new(state),
history: None,
artifacts: None,
metadata: None,
})
};
assert!(ends_stream(&task(TaskState::Completed)));
assert!(ends_stream(&task(TaskState::InputRequired)));
assert!(!ends_stream(&task(TaskState::Working)));
let message = StreamResponse::Message(Message {
id: MessageId::new("m"),
role: MessageRole::Agent,
parts: vec![Part::text("hi")],
task_id: None,
context_id: None,
reference_task_ids: None,
extensions: None,
metadata: None,
});
assert!(ends_stream(&message));
}
#[tokio::test]
async fn a_partial_frame_after_the_final_event_is_not_an_error() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx).with_jsonrpc_envelope(false);
let ev = make_status_event(TaskState::InputRequired, false);
tx.send(Ok(Bytes::from(bare_sse_frame(&ev)))).await.unwrap();
tx.send(Ok(Bytes::from_static(b"data: {\"trunc")))
.await
.unwrap();
drop(tx);
assert!(matches!(stream.next().await, Some(Ok(_))));
assert!(stream.next().await.is_none());
}
#[tokio::test]
async fn last_event_id_persists_across_frames_without_one() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx).with_jsonrpc_envelope(false);
let ev = bare_sse_frame(&make_status_event(TaskState::Working, false));
tx.send(Ok(Bytes::from(format!("id: 4\n{ev}{ev}"))))
.await
.unwrap();
assert!(matches!(stream.next().await, Some(Ok(_))));
assert_eq!(stream.last_event_id(), Some("4"));
assert!(matches!(stream.next().await, Some(Ok(_))));
assert_eq!(stream.last_event_id(), Some("4"), "persists");
}
#[tokio::test]
async fn from_event_channel_delivers_events() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::from_event_channel(rx);
let event = StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
task_id: TaskId::new("task-1"),
context_id: a2a_protocol_types::ContextId::new("ctx-1"),
status: TaskStatus::new(TaskState::Working),
metadata: None,
});
tx.send(Ok(event)).await.expect("send");
drop(tx);
let received = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out")
.expect("a sent event must arrive")
.expect("and must not be an error");
match received {
StreamResponse::StatusUpdate(ev) => {
assert_eq!(ev.task_id, TaskId::new("task-1"));
assert_eq!(ev.status.state, TaskState::Working);
}
other => panic!("expected a status update, got {other:?}"),
}
}
#[tokio::test]
async fn from_event_channel_propagates_errors() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::from_event_channel(rx);
tx.send(Err(ClientError::Transport("frame decode failed".into())))
.await
.expect("send");
drop(tx);
let received = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out")
.expect("an error must be delivered, not swallowed");
assert!(
matches!(received, Err(ClientError::Transport(ref m)) if m == "frame decode failed"),
"the transport's own error must survive the bridge: {received:?}"
);
}
#[tokio::test]
async fn from_event_channel_ends_when_sender_drops() {
let (tx, rx) = mpsc::channel::<ClientResult<StreamResponse>>(8);
let mut stream = EventStream::from_event_channel(rx);
drop(tx);
let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out");
assert!(
matches!(result, Some(Err(ClientError::IncompleteStream { .. }))),
"closing before a final event is incomplete: {result:?}"
);
assert!(stream.next().await.is_none(), "and then the stream ends");
}
#[tokio::test]
async fn from_event_channel_honours_terminal_events() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::from_event_channel(rx);
tx.send(Ok(StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
task_id: TaskId::new("task-1"),
context_id: a2a_protocol_types::ContextId::new("ctx-1"),
status: TaskStatus::new(TaskState::Completed),
metadata: None,
})))
.await
.expect("send");
let first = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out")
.expect("the terminal event itself is delivered");
assert!(first.is_ok());
let next = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out");
assert!(
next.is_none(),
"a terminal event must end the stream even with the sender alive"
);
}
#[tokio::test]
async fn drop_aborts_background_task() {
let (tx, rx) = mpsc::channel::<BodyChunk>(8);
let handle = tokio::spawn(async move {
let _tx = tx;
tokio::time::sleep(Duration::from_secs(60 * 60)).await;
});
let abort_handle = handle.abort_handle();
let stream = EventStream::with_abort_handle(rx, abort_handle);
drop(stream);
let result = tokio::time::timeout(TEST_TIMEOUT, handle)
.await
.expect("timed out waiting for task abort");
assert!(result.is_err(), "task should have been aborted");
assert!(
result.unwrap_err().is_cancelled(),
"task should be cancelled"
);
}
#[test]
fn debug_output_contains_fields() {
let (_tx, rx) = mpsc::channel::<BodyChunk>(8);
let stream = EventStream::new(rx);
let debug = format!("{stream:?}");
assert!(debug.contains("EventStream"), "should contain struct name");
assert!(debug.contains("done: false"), "a fresh stream is not done");
assert!(
debug.contains("pending_frames"),
"should contain 'pending_frames' field"
);
}
#[tokio::test]
async fn debug_output_reports_a_finished_stream_as_done() {
let (tx, rx) = mpsc::channel::<BodyChunk>(8);
let mut stream = EventStream::new(rx);
drop(tx);
let _ = stream.next().await;
assert!(format!("{stream:?}").contains("done: true"));
}
#[test]
fn is_terminal_returns_false_for_working() {
let event = make_status_event(TaskState::Working, false);
assert!(!is_terminal(&event), "Working state should not be terminal");
}
#[test]
fn is_terminal_returns_true_for_completed() {
let event = make_status_event(TaskState::Completed, true);
assert!(is_terminal(&event), "Completed state should be terminal");
}
#[tokio::test]
async fn stream_decodes_jsonrpc_error_as_protocol_error() {
use a2a_protocol_types::{JsonRpcErrorResponse, JsonRpcVersion};
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx);
let error_resp = JsonRpcErrorResponse {
jsonrpc: JsonRpcVersion,
id: Some(serde_json::json!(1)),
error: a2a_protocol_types::JsonRpcError {
code: -32601,
message: "method not found".into(),
data: None,
},
};
let json = serde_json::to_string(&error_resp).unwrap();
let sse_data = format!("data: {json}\n\n");
tx.send(Ok(Bytes::from(sse_data))).await.unwrap();
drop(tx);
let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out")
.unwrap();
assert!(result.is_err(), "JSON-RPC error should produce Err");
match result.unwrap_err() {
ClientError::Protocol(err) => {
assert!(
format!("{err}").contains("method not found"),
"error message should be preserved"
);
}
other => panic!("expected Protocol error, got {other:?}"),
}
let end = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out");
assert!(end.is_none(), "stream should end after JSON-RPC error");
}
#[tokio::test]
async fn stream_invalid_json_returns_serialization_error() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx);
let sse_data = "data: {not valid json}\n\n";
tx.send(Ok(Bytes::from(sse_data))).await.unwrap();
drop(tx);
let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out")
.unwrap();
assert!(result.is_err(), "invalid JSON should produce Err");
assert!(
matches!(result.unwrap_err(), ClientError::Serialization(_)),
"should be a Serialization error"
);
}
#[tokio::test]
async fn stream_drains_parser_after_channel_close() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx);
let event = make_status_event(TaskState::Working, false);
let sse_bytes = sse_frame(&event);
let (first_half, second_half) = sse_bytes.split_at(sse_bytes.len() / 2);
tx.send(Ok(Bytes::from(first_half.to_owned())))
.await
.unwrap();
tx.send(Ok(Bytes::from(second_half.to_owned())))
.await
.unwrap();
drop(tx);
let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out")
.unwrap();
let event = result.unwrap();
assert!(
matches!(event, StreamResponse::StatusUpdate(ref ev) if ev.status.state == TaskState::Working),
"should deliver Working event from drained parser"
);
}
#[tokio::test]
async fn status_code_returns_set_value() {
let (_tx, rx) = mpsc::channel::<BodyChunk>(8);
let stream = EventStream::new(rx);
assert_eq!(stream.status_code(), 200, "default status should be 200");
}
#[tokio::test]
async fn status_code_with_custom_value() {
let (_tx, rx) = mpsc::channel::<BodyChunk>(8);
let task = tokio::spawn(async { tokio::time::sleep(Duration::from_secs(60)).await });
let stream = EventStream::with_status(rx, task.abort_handle(), 201);
assert_eq!(stream.status_code(), 201);
}
#[tokio::test]
async fn first_event_timeout_fires_when_no_data_arrives() {
let (_tx, rx) = mpsc::channel::<BodyChunk>(8);
let mut stream = EventStream::new(rx).with_first_event_timeout(Duration::from_millis(50));
let result = tokio::time::timeout(Duration::from_secs(2), stream.next())
.await
.expect("first-event timeout must fire well within 2s");
assert!(
matches!(result, Some(Err(ClientError::Timeout(ref m))) if m.contains("first-event")),
"expected first-event timeout, got {result:?}"
);
let done = tokio::time::timeout(Duration::from_secs(2), stream.next())
.await
.expect("a completed stream must return promptly");
assert!(done.is_none());
}
#[tokio::test]
async fn first_event_timeout_lifted_after_first_chunk() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx)
.with_jsonrpc_envelope(false)
.with_first_event_timeout(Duration::from_millis(50));
let event = make_status_event(TaskState::Working, false);
tx.send(Ok(Bytes::from(bare_sse_frame(&event))))
.await
.unwrap();
let first = stream.next().await;
assert!(
matches!(first, Some(Ok(_))),
"first event should parse, got {first:?}"
);
let pending = tokio::time::timeout(Duration::from_millis(120), stream.next()).await;
assert!(
pending.is_err(),
"stream must remain open (pending) after first chunk, got {pending:?}"
);
}
const IDLE: Duration = Duration::from_secs(10);
async fn stream_after_one_event() -> (mpsc::Sender<BodyChunk>, EventStream) {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx)
.with_jsonrpc_envelope(false)
.with_idle_timeout(Some(IDLE));
let event = make_status_event(TaskState::Working, false);
tx.send(Ok(Bytes::from(bare_sse_frame(&event))))
.await
.unwrap();
assert!(matches!(stream.next().await, Some(Ok(_))));
(tx, stream)
}
#[tokio::test(start_paused = true)]
async fn idle_timeout_fires_after_the_first_event() {
let (_tx, mut stream) = stream_after_one_event().await;
let start = tokio::time::Instant::now();
let result = stream.next().await;
assert_eq!(start.elapsed(), IDLE, "fires exactly at the bound");
match result {
Some(Err(ClientError::Timeout(msg))) => {
assert!(msg.contains("idle timeout"), "{msg}");
assert!(msg.contains("10s"), "names the bound: {msg}");
}
other => panic!("expected idle Timeout, got {other:?}"),
}
assert!(stream.next().await.is_none(), "the stream has ended");
}
#[tokio::test(start_paused = true)]
async fn cancelled_polls_do_not_restart_the_idle_clock() {
let (_tx, mut stream) = stream_after_one_event().await;
for _ in 0..3 {
let poll = tokio::time::timeout(IDLE / 4, stream.next()).await;
assert!(poll.is_err(), "still inside the bound: {poll:?}");
}
let start = tokio::time::Instant::now();
let result = tokio::time::timeout(IDLE, stream.next())
.await
.expect("the original deadline falls inside this poll");
assert_eq!(start.elapsed(), IDLE / 4);
assert!(matches!(result, Some(Err(ClientError::Timeout(_)))));
}
#[tokio::test(start_paused = true)]
async fn a_keep_alive_comment_restarts_the_idle_clock() {
let (tx, mut stream) = stream_after_one_event().await;
let quiet = IDLE * 3 / 4;
tokio::time::advance(quiet).await;
tx.send(Ok(Bytes::from_static(b": keep-alive\n\n")))
.await
.unwrap();
let poll = tokio::time::timeout(quiet, stream.next()).await;
assert!(poll.is_err(), "the comment renewed the bound: {poll:?}");
let result = stream.next().await;
assert!(matches!(result, Some(Err(ClientError::Timeout(_)))));
}
#[tokio::test(start_paused = true)]
async fn an_absent_or_unrepresentable_idle_bound_never_fires() {
for bound in [None, Some(Duration::MAX)] {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx)
.with_jsonrpc_envelope(false)
.with_idle_timeout(bound);
let event = make_status_event(TaskState::Working, false);
tx.send(Ok(Bytes::from(bare_sse_frame(&event))))
.await
.unwrap();
assert!(matches!(stream.next().await, Some(Ok(_))));
let poll = tokio::time::timeout(IDLE * 1000, stream.next()).await;
assert!(poll.is_err(), "{bound:?} must not fire: {poll:?}");
}
}
#[tokio::test(start_paused = true)]
async fn streams_carry_the_default_idle_bound() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx).with_jsonrpc_envelope(false);
let event = make_status_event(TaskState::Working, false);
tx.send(Ok(Bytes::from(bare_sse_frame(&event))))
.await
.unwrap();
assert!(matches!(stream.next().await, Some(Ok(_))));
let start = tokio::time::Instant::now();
let result = stream.next().await;
assert_eq!(
start.elapsed(),
crate::config::DEFAULT_STREAM_IDLE_TIMEOUT,
"the default bound applies"
);
assert!(matches!(result, Some(Err(ClientError::Timeout(_)))));
}
#[tokio::test(start_paused = true)]
async fn idle_expiry_aborts_the_body_reader() {
let (tx, rx) = mpsc::channel::<BodyChunk>(8);
let reader = tokio::spawn(async move {
let event = make_status_event(TaskState::Working, false);
let _ = tx.send(Ok(Bytes::from(bare_sse_frame(&event)))).await;
std::future::pending::<()>().await;
});
let mut stream = EventStream::with_abort_handle(rx, reader.abort_handle())
.with_jsonrpc_envelope(false)
.with_idle_timeout(Some(IDLE));
assert!(matches!(stream.next().await, Some(Ok(_))));
assert!(matches!(
stream.next().await,
Some(Err(ClientError::Timeout(_)))
));
let joined = reader.await;
assert!(
joined.is_err_and(|e| e.is_cancelled()),
"the reader must be aborted at expiry while the stream is still alive"
);
drop(stream);
}
#[tokio::test]
async fn stream_transport_error_from_channel() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx);
tx.send(Err(ClientError::HttpClient("connection reset".into())))
.await
.unwrap();
let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out")
.unwrap();
match result {
Err(ClientError::HttpClient(msg)) => {
assert!(msg.contains("connection reset"));
}
other => panic!("expected HttpClient error, got {other:?}"),
}
let end = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out");
assert!(end.is_none(), "stream should end after transport error");
}
#[tokio::test]
async fn non_terminal_event_does_not_end_stream() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx);
let working = make_status_event(TaskState::Working, false);
let completed = make_status_event(TaskState::Completed, true);
tx.send(Ok(Bytes::from(sse_frame(&working)))).await.unwrap();
tx.send(Ok(Bytes::from(sse_frame(&completed))))
.await
.unwrap();
let first = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out on first event")
.unwrap()
.unwrap();
assert!(
matches!(first, StreamResponse::StatusUpdate(ref ev) if ev.status.state == TaskState::Working)
);
let second = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out on second event")
.unwrap()
.unwrap();
assert!(
matches!(second, StreamResponse::StatusUpdate(ref ev) if ev.status.state == TaskState::Completed)
);
let end = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out waiting for stream end");
assert!(end.is_none());
}
fn bare_sse_frame(event: &StreamResponse) -> String {
let json = serde_json::to_string(event).unwrap();
format!("data: {json}\n\n")
}
#[tokio::test]
async fn bare_stream_delivers_events() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx).with_jsonrpc_envelope(false);
let event = make_status_event(TaskState::Working, false);
tx.send(Ok(Bytes::from(bare_sse_frame(&event))))
.await
.unwrap();
drop(tx);
let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out")
.unwrap()
.unwrap();
assert!(
matches!(result, StreamResponse::StatusUpdate(ref ev) if ev.status.state == TaskState::Working)
);
}
#[tokio::test]
async fn bare_stream_ends_on_terminal() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx).with_jsonrpc_envelope(false);
let event = make_status_event(TaskState::Completed, true);
tx.send(Ok(Bytes::from(bare_sse_frame(&event))))
.await
.unwrap();
let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out")
.unwrap()
.unwrap();
assert!(
matches!(result, StreamResponse::StatusUpdate(ref ev) if ev.status.state == TaskState::Completed)
);
let end = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out");
assert!(end.is_none(), "bare stream should end after terminal event");
}
#[tokio::test]
async fn bare_stream_rejects_jsonrpc_envelope() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx).with_jsonrpc_envelope(false);
let event = make_status_event(TaskState::Working, false);
let envelope_frame = sse_frame(&event); tx.send(Ok(Bytes::from(envelope_frame))).await.unwrap();
drop(tx);
let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out")
.unwrap();
assert!(
result.is_err(),
"bare stream should reject JSON-RPC envelope as invalid"
);
}
#[tokio::test]
async fn envelope_stream_rejects_bare_response() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx);
let event = make_status_event(TaskState::Working, false);
let bare_frame = bare_sse_frame(&event);
tx.send(Ok(Bytes::from(bare_frame))).await.unwrap();
drop(tx);
let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out")
.unwrap();
assert!(
result.is_err(),
"envelope stream should reject bare StreamResponse"
);
}
#[tokio::test]
async fn bare_stream_multiple_events() {
let (tx, rx) = mpsc::channel(8);
let mut stream = EventStream::new(rx).with_jsonrpc_envelope(false);
let working = make_status_event(TaskState::Working, false);
let completed = make_status_event(TaskState::Completed, true);
tx.send(Ok(Bytes::from(bare_sse_frame(&working))))
.await
.unwrap();
tx.send(Ok(Bytes::from(bare_sse_frame(&completed))))
.await
.unwrap();
let first = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out")
.unwrap()
.unwrap();
assert!(
matches!(first, StreamResponse::StatusUpdate(ref ev) if ev.status.state == TaskState::Working)
);
let second = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out")
.unwrap()
.unwrap();
assert!(
matches!(second, StreamResponse::StatusUpdate(ref ev) if ev.status.state == TaskState::Completed)
);
let end = tokio::time::timeout(TEST_TIMEOUT, stream.next())
.await
.expect("timed out");
assert!(end.is_none());
}
}