1use a2a_protocol_types::{JsonRpcResponse, StreamResponse};
42use hyper::body::Bytes;
43use tokio::sync::mpsc;
44use tokio::task::AbortHandle;
45
46use crate::error::{ClientError, ClientResult};
47use crate::streaming::sse_parser::SseParser;
48
49pub(crate) type BodyChunk = ClientResult<Bytes>;
53
54const EVENT_BRIDGE_CAPACITY: usize = 64;
60
61pub struct EventStream {
72 rx: mpsc::Receiver<BodyChunk>,
74 parser: SseParser,
76 phase: Phase,
78 abort_handle: Option<AbortHandle>,
80 status_code: u16,
87 jsonrpc_envelope: bool,
93 first_event_timeout: Option<std::time::Duration>,
101 idle_timeout: Option<std::time::Duration>,
104 last_activity: tokio::time::Instant,
109 first_chunk_received: bool,
112 at_final_event: bool,
116 events_received: u64,
118 last_event_id: Option<String>,
120 held: Option<
133 Box<dyn std::any::Any + Send + Sync + std::panic::UnwindSafe + std::panic::RefUnwindSafe>,
134 >,
135}
136
137impl EventStream {
138 #[must_use]
144 #[cfg(any(test, feature = "websocket"))]
145 pub(crate) fn new(rx: mpsc::Receiver<BodyChunk>) -> Self {
146 Self {
147 rx,
148 parser: SseParser::new(),
149 phase: Phase::Open,
150 abort_handle: None,
151 status_code: 200,
152 jsonrpc_envelope: true,
153 first_event_timeout: None,
154 idle_timeout: Some(crate::config::DEFAULT_STREAM_IDLE_TIMEOUT),
155 last_activity: tokio::time::Instant::now(),
156 first_chunk_received: false,
157 at_final_event: false,
158 events_received: 0,
159 last_event_id: None,
160 held: None,
161 }
162 }
163
164 #[must_use]
169 #[cfg(test)]
170 pub(crate) fn with_abort_handle(
171 rx: mpsc::Receiver<BodyChunk>,
172 abort_handle: AbortHandle,
173 ) -> Self {
174 Self {
175 rx,
176 parser: SseParser::new(),
177 phase: Phase::Open,
178 abort_handle: Some(abort_handle),
179 status_code: 200,
180 jsonrpc_envelope: true,
181 first_event_timeout: None,
182 idle_timeout: Some(crate::config::DEFAULT_STREAM_IDLE_TIMEOUT),
183 last_activity: tokio::time::Instant::now(),
184 first_chunk_received: false,
185 at_final_event: false,
186 events_received: 0,
187 last_event_id: None,
188 held: None,
189 }
190 }
191
192 #[must_use]
229 pub fn from_event_channel(mut rx: mpsc::Receiver<ClientResult<StreamResponse>>) -> Self {
230 let (tx, body_rx) = mpsc::channel::<BodyChunk>(EVENT_BRIDGE_CAPACITY);
231
232 let bridge = tokio::spawn(async move {
238 while let Some(event) = rx.recv().await {
239 let chunk = match event {
240 Ok(ref ev) => serde_json::to_string(ev).map_or_else(
241 |e| Err(ClientError::Serialization(e)),
242 |json| {
243 Ok(Bytes::from(format!(
244 "data: {{\"jsonrpc\":\"2.0\",\"id\":null,\"result\":{json}}}\n\n"
245 )))
246 },
247 ),
248 Err(e) => Err(e),
249 };
250 if tx.send(chunk).await.is_err() {
251 break;
252 }
253 }
254 });
255
256 Self::with_status(body_rx, bridge.abort_handle(), 200)
257 }
258
259 #[must_use]
262 pub(crate) fn with_status(
263 rx: mpsc::Receiver<BodyChunk>,
264 abort_handle: AbortHandle,
265 status_code: u16,
266 ) -> Self {
267 Self {
268 rx,
269 parser: SseParser::new(),
270 phase: Phase::Open,
271 abort_handle: Some(abort_handle),
272 status_code,
273 jsonrpc_envelope: true,
274 first_event_timeout: None,
275 idle_timeout: Some(crate::config::DEFAULT_STREAM_IDLE_TIMEOUT),
276 last_activity: tokio::time::Instant::now(),
277 first_chunk_received: false,
278 at_final_event: false,
279 events_received: 0,
280 last_event_id: None,
281 held: None,
282 }
283 }
284
285 #[must_use]
290 pub(crate) const fn with_jsonrpc_envelope(mut self, envelope: bool) -> Self {
291 self.jsonrpc_envelope = envelope;
292 self
293 }
294
295 #[must_use]
310 pub const fn with_first_event_timeout(mut self, timeout: std::time::Duration) -> Self {
311 self.first_event_timeout = Some(timeout);
312 self
313 }
314
315 #[must_use]
330 pub const fn with_idle_timeout(mut self, timeout: Option<std::time::Duration>) -> Self {
331 self.idle_timeout = timeout;
332 self
333 }
334
335 #[must_use]
341 pub const fn with_max_event_size(mut self, max_bytes: usize) -> Self {
342 self.parser.set_max_event_size(max_bytes);
343 self
344 }
345
346 #[must_use]
361 #[cfg(feature = "websocket")]
362 pub(crate) fn holding(
363 mut self,
364 resource: impl std::any::Any
365 + Send
366 + Sync
367 + std::panic::UnwindSafe
368 + std::panic::RefUnwindSafe
369 + 'static,
370 ) -> Self {
371 self.held = Some(Box::new(resource));
372 self
373 }
374
375 #[must_use]
380 pub const fn status_code(&self) -> u16 {
381 self.status_code
382 }
383
384 #[must_use]
395 pub fn last_event_id(&self) -> Option<&str> {
396 self.last_event_id.as_deref()
397 }
398
399 pub async fn next(&mut self) -> Option<ClientResult<StreamResponse>> {
411 loop {
412 if let Some(result) = self.parser.next_frame() {
414 match result {
415 Ok(frame) => {
416 self.last_event_id = frame.id;
417 return Some(self.decode_frame(&frame.data));
418 }
419 Err(e) => {
420 return Some(Err(ClientError::Transport(e.to_string())));
421 }
422 }
423 }
424
425 match self.phase {
426 Phase::Finished => return None,
427 Phase::BodyEnded => {
428 self.phase = Phase::Finished;
429 return self.judge_end().map(Err);
430 }
431 Phase::Open => {}
432 }
433
434 let chunk = match self.silence_deadline() {
437 Some((deadline, bound)) => {
438 let Ok(chunk) = tokio::time::timeout_at(deadline, self.rx.recv()).await else {
439 return Some(Err(self.fail_silent(bound)));
440 };
441 chunk
442 }
443 None => self.rx.recv().await,
444 };
445 match chunk {
446 None => {
447 self.phase = Phase::BodyEnded;
451 }
452 Some(Err(e)) => {
453 self.phase = Phase::Finished;
454 return Some(Err(e));
455 }
456 Some(Ok(bytes)) => {
457 self.first_chunk_received = true;
458 self.last_activity = tokio::time::Instant::now();
459 self.parser.feed(&bytes);
460 }
461 }
462 }
463 }
464
465 fn silence_deadline(&self) -> Option<(tokio::time::Instant, std::time::Duration)> {
474 let bound = if self.first_chunk_received {
475 self.idle_timeout
476 } else {
477 self.first_event_timeout
478 }?;
479 Some((self.last_activity.checked_add(bound)?, bound))
480 }
481
482 fn fail_silent(&mut self, bound: std::time::Duration) -> ClientError {
487 self.phase = Phase::Finished;
488 if let Some(handle) = self.abort_handle.take() {
489 handle.abort();
490 }
491 if self.first_chunk_received {
492 ClientError::Timeout(format!(
493 "stream idle timeout: no data, not even a keep-alive comment, for {bound:?} \
494 since the last data received; the task on the server is not cancelled, \
495 so resubscribe to continue"
496 ))
497 } else {
498 ClientError::Timeout(format!(
499 "stream produced no data before the first-event timeout ({bound:?})"
500 ))
501 }
502 }
503
504 fn judge_end(&self) -> Option<ClientError> {
507 let partial = self.parser.incomplete_event_len();
508 if self.at_final_event {
509 if let Some(bytes) = partial {
510 trace_warn!(
513 bytes,
514 "SSE body ended mid-frame after the stream's final event; \
515 discarding the unterminated frame"
516 );
517 let _ = bytes;
518 }
519 return None;
520 }
521 let detail = match (partial, self.events_received) {
522 (Some(bytes), _) => format!(
523 "the connection closed mid-frame, discarding {bytes} bytes of an \
524 unterminated event"
525 ),
526 (None, 0) => "the body ended before any event".to_owned(),
527 (None, n) => format!(
528 "the body ended after {n} event(s), none of them a final one \
529 (a Message, or a terminal or interrupted task state)"
530 ),
531 };
532 Some(ClientError::IncompleteStream {
533 last_event_id: self.last_event_id.clone(),
534 detail,
535 })
536 }
537
538 const fn record(&mut self, event: &StreamResponse) {
540 self.events_received += 1;
541 self.at_final_event = ends_stream(event);
542 }
543
544 fn decode_frame(&mut self, data: &str) -> ClientResult<StreamResponse> {
545 if self.jsonrpc_envelope {
546 let envelope: JsonRpcResponse<StreamResponse> =
548 serde_json::from_str(data).map_err(ClientError::Serialization)?;
549
550 match envelope {
551 JsonRpcResponse::Success(ok) => {
552 self.record(&ok.result);
553 if is_terminal(&ok.result) {
554 self.phase = Phase::Finished;
555 }
556 Ok(ok.result)
557 }
558 JsonRpcResponse::Error(err) => {
559 self.phase = Phase::Finished;
560 let a2a = crate::transport::map_jsonrpc_error(
561 err.error.code,
562 err.error.message,
563 err.error.data,
564 );
565 Err(ClientError::Protocol(a2a))
566 }
567 }
568 } else {
569 let event: StreamResponse = match serde_json::from_str(data) {
574 Ok(event) => event,
575 Err(e) => {
576 return Err(
577 match crate::transport::rest::error_frame::decode_stream_error_frame(data) {
578 Some(a2a) => {
579 self.phase = Phase::Finished;
580 ClientError::Protocol(a2a)
581 }
582 None => ClientError::Serialization(e),
583 },
584 );
585 }
586 };
587 self.record(&event);
588 if is_terminal(&event) {
589 self.phase = Phase::Finished;
590 }
591 Ok(event)
592 }
593 }
594}
595
596impl Drop for EventStream {
597 fn drop(&mut self) {
598 if let Some(handle) = self.abort_handle.take() {
599 handle.abort();
600 }
601 drop(self.held.take());
608 }
609}
610
611#[allow(clippy::missing_fields_in_debug)]
612impl std::fmt::Debug for EventStream {
613 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
614 f.debug_struct("EventStream")
616 .field("done", &(self.phase == Phase::Finished))
617 .field("pending_frames", &self.parser.pending_count())
618 .finish()
619 }
620}
621
622#[derive(Debug, Clone, Copy, PartialEq, Eq)]
624enum Phase {
625 Open,
627 BodyEnded,
630 Finished,
632}
633
634const fn ends_stream(event: &StreamResponse) -> bool {
645 let state = match event {
646 StreamResponse::Message(_) => return true,
647 StreamResponse::Task(task) => task.status.state,
648 StreamResponse::StatusUpdate(update) => update.status.state,
649 _ => return false,
650 };
651 state.is_terminal() || state.is_interrupted()
652}
653
654const fn is_terminal(event: &StreamResponse) -> bool {
656 matches!(
657 event,
658 StreamResponse::StatusUpdate(ev) if ev.status.state.is_terminal()
659 )
660}
661
662#[cfg(test)]
676const fn _event_stream_is_unwind_safe() {
677 const fn assert_unwind_safe<T: std::panic::UnwindSafe + std::panic::RefUnwindSafe>() {}
678 assert_unwind_safe::<EventStream>();
679}
680
681#[cfg(test)]
682mod tests {
683 use super::*;
684 use a2a_protocol_types::{
685 JsonRpcSuccessResponse, JsonRpcVersion, TaskId, TaskState, TaskStatus,
686 TaskStatusUpdateEvent,
687 };
688 use std::time::Duration;
689
690 const TEST_TIMEOUT: Duration = Duration::from_secs(5);
693
694 fn make_status_event(state: TaskState, _is_final: bool) -> StreamResponse {
695 StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
696 task_id: TaskId::new("t1"),
697 context_id: a2a_protocol_types::ContextId::new("c1"),
698 status: TaskStatus {
699 state,
700 message: None,
701 timestamp: None,
702 },
703 metadata: None,
704 })
705 }
706
707 fn sse_frame(event: &StreamResponse) -> String {
708 let resp = JsonRpcSuccessResponse {
709 jsonrpc: JsonRpcVersion,
710 id: Some(serde_json::json!(1)),
711 result: event.clone(),
712 };
713 let json = serde_json::to_string(&resp).unwrap();
714 format!("data: {json}\n\n")
715 }
716
717 #[tokio::test]
718 async fn stream_delivers_events() {
719 let (tx, rx) = mpsc::channel(8);
720 let mut stream = EventStream::new(rx);
721
722 let event = make_status_event(TaskState::Working, false);
723 let sse_bytes = sse_frame(&event);
724 tx.send(Ok(Bytes::from(sse_bytes))).await.unwrap();
725 drop(tx);
726
727 let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
728 .await
729 .expect("timed out")
730 .unwrap()
731 .unwrap();
732 assert!(
733 matches!(result, StreamResponse::StatusUpdate(ref ev) if ev.status.state == TaskState::Working)
734 );
735 }
736
737 #[tokio::test]
738 async fn stream_ends_on_final_event() {
739 let (tx, rx) = mpsc::channel(8);
740 let mut stream = EventStream::new(rx);
741
742 let event = make_status_event(TaskState::Completed, true);
743 let sse_bytes = sse_frame(&event);
744 tx.send(Ok(Bytes::from(sse_bytes))).await.unwrap();
745
746 let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
748 .await
749 .expect("timed out waiting for final event")
750 .unwrap()
751 .unwrap();
752 assert!(
753 matches!(result, StreamResponse::StatusUpdate(ref ev) if ev.status.state == TaskState::Completed)
754 );
755
756 let end = tokio::time::timeout(TEST_TIMEOUT, stream.next())
758 .await
759 .expect("timed out waiting for stream end");
760 assert!(end.is_none());
761 }
762
763 #[tokio::test]
764 async fn stream_propagates_body_error() {
765 let (tx, rx) = mpsc::channel(8);
766 let mut stream = EventStream::new(rx);
767
768 tx.send(Err(ClientError::Transport("network error".into())))
769 .await
770 .unwrap();
771
772 let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
773 .await
774 .expect("timed out")
775 .unwrap();
776 assert!(result.is_err());
777 }
778
779 #[tokio::test]
782 async fn a_channel_closed_before_any_event_is_incomplete() {
783 let (tx, rx) = mpsc::channel(8);
784 let mut stream = EventStream::new(rx);
785 drop(tx);
786
787 let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
788 .await
789 .expect("timed out");
790 match result {
791 Some(Err(ClientError::IncompleteStream {
792 last_event_id,
793 detail,
794 })) => {
795 assert_eq!(last_event_id, None);
796 assert_eq!(detail, "the body ended before any event");
797 }
798 other => panic!("expected IncompleteStream, got {other:?}"),
799 }
800 assert!(
801 stream.next().await.is_none(),
802 "reported once, then finished"
803 );
804 }
805
806 #[tokio::test]
809 async fn the_end_of_the_body_is_judged_by_the_last_event() {
810 for (states, clean) in [
811 (vec![TaskState::Working, TaskState::InputRequired], true),
812 (vec![TaskState::Working, TaskState::AuthRequired], true),
813 (vec![TaskState::InputRequired, TaskState::Working], false),
814 (vec![TaskState::Submitted, TaskState::Working], false),
815 ] {
816 let (tx, rx) = mpsc::channel(8);
817 let mut stream = EventStream::new(rx).with_jsonrpc_envelope(false);
818 for state in &states {
819 let ev = make_status_event(*state, false);
820 tx.send(Ok(Bytes::from(bare_sse_frame(&ev)))).await.unwrap();
821 }
822 drop(tx);
823 for _ in &states {
824 assert!(matches!(stream.next().await, Some(Ok(_))), "{states:?}");
825 }
826 let end = stream.next().await;
827 if clean {
828 assert!(end.is_none(), "{states:?}: a completion, got {end:?}");
829 } else {
830 match end {
831 Some(Err(ClientError::IncompleteStream { detail, .. })) => {
832 assert!(detail.contains("after 2 event(s)"), "{detail}");
833 }
834 other => panic!("{states:?}: expected IncompleteStream, got {other:?}"),
835 }
836 }
837 assert!(stream.next().await.is_none());
838 }
839 }
840
841 #[test]
844 fn ends_stream_covers_final_events_only() {
845 use a2a_protocol_types::{Message, MessageId, MessageRole, Part, Task};
846 for state in [
847 TaskState::Completed,
848 TaskState::Failed,
849 TaskState::Canceled,
850 TaskState::Rejected,
851 TaskState::InputRequired,
852 TaskState::AuthRequired,
853 ] {
854 assert!(ends_stream(&make_status_event(state, false)), "{state:?}");
855 }
856 for state in [TaskState::Submitted, TaskState::Working] {
857 assert!(!ends_stream(&make_status_event(state, false)), "{state:?}");
858 }
859 let task = |state| {
860 StreamResponse::Task(Task {
861 id: TaskId::new("t"),
862 context_id: a2a_protocol_types::ContextId::new("c"),
863 status: TaskStatus::new(state),
864 history: None,
865 artifacts: None,
866 metadata: None,
867 })
868 };
869 assert!(ends_stream(&task(TaskState::Completed)));
870 assert!(ends_stream(&task(TaskState::InputRequired)));
871 assert!(!ends_stream(&task(TaskState::Working)));
872 let message = StreamResponse::Message(Message {
873 id: MessageId::new("m"),
874 role: MessageRole::Agent,
875 parts: vec![Part::text("hi")],
876 task_id: None,
877 context_id: None,
878 reference_task_ids: None,
879 extensions: None,
880 metadata: None,
881 });
882 assert!(ends_stream(&message));
883 }
884
885 #[tokio::test]
888 async fn a_partial_frame_after_the_final_event_is_not_an_error() {
889 let (tx, rx) = mpsc::channel(8);
890 let mut stream = EventStream::new(rx).with_jsonrpc_envelope(false);
891 let ev = make_status_event(TaskState::InputRequired, false);
892 tx.send(Ok(Bytes::from(bare_sse_frame(&ev)))).await.unwrap();
893 tx.send(Ok(Bytes::from_static(b"data: {\"trunc")))
894 .await
895 .unwrap();
896 drop(tx);
897 assert!(matches!(stream.next().await, Some(Ok(_))));
898 assert!(stream.next().await.is_none());
899 }
900
901 #[tokio::test]
904 async fn last_event_id_persists_across_frames_without_one() {
905 let (tx, rx) = mpsc::channel(8);
906 let mut stream = EventStream::new(rx).with_jsonrpc_envelope(false);
907 let ev = bare_sse_frame(&make_status_event(TaskState::Working, false));
908 tx.send(Ok(Bytes::from(format!("id: 4\n{ev}{ev}"))))
909 .await
910 .unwrap();
911 assert!(matches!(stream.next().await, Some(Ok(_))));
912 assert_eq!(stream.last_event_id(), Some("4"));
913 assert!(matches!(stream.next().await, Some(Ok(_))));
914 assert_eq!(stream.last_event_id(), Some("4"), "persists");
915 }
916
917 #[tokio::test]
925 async fn from_event_channel_delivers_events() {
926 let (tx, rx) = mpsc::channel(8);
927 let mut stream = EventStream::from_event_channel(rx);
928
929 let event = StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
930 task_id: TaskId::new("task-1"),
931 context_id: a2a_protocol_types::ContextId::new("ctx-1"),
932 status: TaskStatus::new(TaskState::Working),
933 metadata: None,
934 });
935 tx.send(Ok(event)).await.expect("send");
936 drop(tx);
937
938 let received = tokio::time::timeout(TEST_TIMEOUT, stream.next())
939 .await
940 .expect("timed out")
941 .expect("a sent event must arrive")
942 .expect("and must not be an error");
943
944 match received {
945 StreamResponse::StatusUpdate(ev) => {
946 assert_eq!(ev.task_id, TaskId::new("task-1"));
947 assert_eq!(ev.status.state, TaskState::Working);
948 }
949 other => panic!("expected a status update, got {other:?}"),
950 }
951 }
952
953 #[tokio::test]
957 async fn from_event_channel_propagates_errors() {
958 let (tx, rx) = mpsc::channel(8);
959 let mut stream = EventStream::from_event_channel(rx);
960
961 tx.send(Err(ClientError::Transport("frame decode failed".into())))
962 .await
963 .expect("send");
964 drop(tx);
965
966 let received = tokio::time::timeout(TEST_TIMEOUT, stream.next())
967 .await
968 .expect("timed out")
969 .expect("an error must be delivered, not swallowed");
970
971 assert!(
972 matches!(received, Err(ClientError::Transport(ref m)) if m == "frame decode failed"),
973 "the transport's own error must survive the bridge: {received:?}"
974 );
975 }
976
977 #[tokio::test]
980 async fn from_event_channel_ends_when_sender_drops() {
981 let (tx, rx) = mpsc::channel::<ClientResult<StreamResponse>>(8);
982 let mut stream = EventStream::from_event_channel(rx);
983 drop(tx);
984
985 let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
986 .await
987 .expect("timed out");
988
989 assert!(
990 matches!(result, Some(Err(ClientError::IncompleteStream { .. }))),
991 "closing before a final event is incomplete: {result:?}"
992 );
993 assert!(stream.next().await.is_none(), "and then the stream ends");
994 }
995
996 #[tokio::test]
1000 async fn from_event_channel_honours_terminal_events() {
1001 let (tx, rx) = mpsc::channel(8);
1002 let mut stream = EventStream::from_event_channel(rx);
1003
1004 tx.send(Ok(StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
1005 task_id: TaskId::new("task-1"),
1006 context_id: a2a_protocol_types::ContextId::new("ctx-1"),
1007 status: TaskStatus::new(TaskState::Completed),
1008 metadata: None,
1009 })))
1010 .await
1011 .expect("send");
1012
1013 let first = tokio::time::timeout(TEST_TIMEOUT, stream.next())
1014 .await
1015 .expect("timed out")
1016 .expect("the terminal event itself is delivered");
1017 assert!(first.is_ok());
1018
1019 let next = tokio::time::timeout(TEST_TIMEOUT, stream.next())
1022 .await
1023 .expect("timed out");
1024 assert!(
1025 next.is_none(),
1026 "a terminal event must end the stream even with the sender alive"
1027 );
1028 }
1029
1030 #[tokio::test]
1031 async fn drop_aborts_background_task() {
1032 let (tx, rx) = mpsc::channel::<BodyChunk>(8);
1033 let handle = tokio::spawn(async move {
1035 let _tx = tx;
1037 tokio::time::sleep(Duration::from_secs(60 * 60)).await;
1039 });
1040 let abort_handle = handle.abort_handle();
1041 let stream = EventStream::with_abort_handle(rx, abort_handle);
1042 drop(stream);
1044 let result = tokio::time::timeout(TEST_TIMEOUT, handle)
1046 .await
1047 .expect("timed out waiting for task abort");
1048 assert!(result.is_err(), "task should have been aborted");
1049 assert!(
1050 result.unwrap_err().is_cancelled(),
1051 "task should be cancelled"
1052 );
1053 }
1054
1055 #[test]
1056 fn debug_output_contains_fields() {
1057 let (_tx, rx) = mpsc::channel::<BodyChunk>(8);
1058 let stream = EventStream::new(rx);
1059 let debug = format!("{stream:?}");
1060 assert!(debug.contains("EventStream"), "should contain struct name");
1061 assert!(debug.contains("done: false"), "a fresh stream is not done");
1062 assert!(
1063 debug.contains("pending_frames"),
1064 "should contain 'pending_frames' field"
1065 );
1066 }
1067
1068 #[tokio::test]
1070 async fn debug_output_reports_a_finished_stream_as_done() {
1071 let (tx, rx) = mpsc::channel::<BodyChunk>(8);
1072 let mut stream = EventStream::new(rx);
1073 drop(tx);
1074 let _ = stream.next().await;
1075 assert!(format!("{stream:?}").contains("done: true"));
1076 }
1077
1078 #[test]
1079 fn is_terminal_returns_false_for_working() {
1080 let event = make_status_event(TaskState::Working, false);
1081 assert!(!is_terminal(&event), "Working state should not be terminal");
1082 }
1083
1084 #[test]
1085 fn is_terminal_returns_true_for_completed() {
1086 let event = make_status_event(TaskState::Completed, true);
1087 assert!(is_terminal(&event), "Completed state should be terminal");
1088 }
1089
1090 #[tokio::test]
1093 async fn stream_decodes_jsonrpc_error_as_protocol_error() {
1094 use a2a_protocol_types::{JsonRpcErrorResponse, JsonRpcVersion};
1095
1096 let (tx, rx) = mpsc::channel(8);
1097 let mut stream = EventStream::new(rx);
1098
1099 let error_resp = JsonRpcErrorResponse {
1101 jsonrpc: JsonRpcVersion,
1102 id: Some(serde_json::json!(1)),
1103 error: a2a_protocol_types::JsonRpcError {
1104 code: -32601,
1105 message: "method not found".into(),
1106 data: None,
1107 },
1108 };
1109 let json = serde_json::to_string(&error_resp).unwrap();
1110 let sse_data = format!("data: {json}\n\n");
1111 tx.send(Ok(Bytes::from(sse_data))).await.unwrap();
1112 drop(tx);
1113
1114 let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
1115 .await
1116 .expect("timed out")
1117 .unwrap();
1118 assert!(result.is_err(), "JSON-RPC error should produce Err");
1119 match result.unwrap_err() {
1120 ClientError::Protocol(err) => {
1121 assert!(
1122 format!("{err}").contains("method not found"),
1123 "error message should be preserved"
1124 );
1125 }
1126 other => panic!("expected Protocol error, got {other:?}"),
1127 }
1128
1129 let end = tokio::time::timeout(TEST_TIMEOUT, stream.next())
1131 .await
1132 .expect("timed out");
1133 assert!(end.is_none(), "stream should end after JSON-RPC error");
1134 }
1135
1136 #[tokio::test]
1139 async fn stream_invalid_json_returns_serialization_error() {
1140 let (tx, rx) = mpsc::channel(8);
1141 let mut stream = EventStream::new(rx);
1142
1143 let sse_data = "data: {not valid json}\n\n";
1144 tx.send(Ok(Bytes::from(sse_data))).await.unwrap();
1145 drop(tx);
1146
1147 let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
1148 .await
1149 .expect("timed out")
1150 .unwrap();
1151 assert!(result.is_err(), "invalid JSON should produce Err");
1152 assert!(
1153 matches!(result.unwrap_err(), ClientError::Serialization(_)),
1154 "should be a Serialization error"
1155 );
1156 }
1157
1158 #[tokio::test]
1161 async fn stream_drains_parser_after_channel_close() {
1162 let (tx, rx) = mpsc::channel(8);
1163 let mut stream = EventStream::new(rx);
1164
1165 let event = make_status_event(TaskState::Working, false);
1168 let sse_bytes = sse_frame(&event);
1169 let (first_half, second_half) = sse_bytes.split_at(sse_bytes.len() / 2);
1170
1171 tx.send(Ok(Bytes::from(first_half.to_owned())))
1172 .await
1173 .unwrap();
1174 tx.send(Ok(Bytes::from(second_half.to_owned())))
1175 .await
1176 .unwrap();
1177 drop(tx);
1178
1179 let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
1180 .await
1181 .expect("timed out")
1182 .unwrap();
1183 let event = result.unwrap();
1184 assert!(
1185 matches!(event, StreamResponse::StatusUpdate(ref ev) if ev.status.state == TaskState::Working),
1186 "should deliver Working event from drained parser"
1187 );
1188 }
1189
1190 #[tokio::test]
1192 async fn status_code_returns_set_value() {
1193 let (_tx, rx) = mpsc::channel::<BodyChunk>(8);
1194 let stream = EventStream::new(rx);
1195 assert_eq!(stream.status_code(), 200, "default status should be 200");
1196 }
1197
1198 #[tokio::test]
1200 async fn status_code_with_custom_value() {
1201 let (_tx, rx) = mpsc::channel::<BodyChunk>(8);
1202 let task = tokio::spawn(async { tokio::time::sleep(Duration::from_secs(60)).await });
1203 let stream = EventStream::with_status(rx, task.abort_handle(), 201);
1204 assert_eq!(stream.status_code(), 201);
1205 }
1206
1207 #[tokio::test]
1210 async fn first_event_timeout_fires_when_no_data_arrives() {
1211 let (_tx, rx) = mpsc::channel::<BodyChunk>(8);
1213 let mut stream = EventStream::new(rx).with_first_event_timeout(Duration::from_millis(50));
1214 let result = tokio::time::timeout(Duration::from_secs(2), stream.next())
1217 .await
1218 .expect("first-event timeout must fire well within 2s");
1219 assert!(
1220 matches!(result, Some(Err(ClientError::Timeout(ref m))) if m.contains("first-event")),
1221 "expected first-event timeout, got {result:?}"
1222 );
1223 let done = tokio::time::timeout(Duration::from_secs(2), stream.next())
1225 .await
1226 .expect("a completed stream must return promptly");
1227 assert!(done.is_none());
1228 }
1229
1230 #[tokio::test]
1233 async fn first_event_timeout_lifted_after_first_chunk() {
1234 let (tx, rx) = mpsc::channel(8);
1235 let mut stream = EventStream::new(rx)
1236 .with_jsonrpc_envelope(false)
1237 .with_first_event_timeout(Duration::from_millis(50));
1238 let event = make_status_event(TaskState::Working, false);
1240 tx.send(Ok(Bytes::from(bare_sse_frame(&event))))
1241 .await
1242 .unwrap();
1243 let first = stream.next().await;
1244 assert!(
1245 matches!(first, Some(Ok(_))),
1246 "first event should parse, got {first:?}"
1247 );
1248 let pending = tokio::time::timeout(Duration::from_millis(120), stream.next()).await;
1251 assert!(
1252 pending.is_err(),
1253 "stream must remain open (pending) after first chunk, got {pending:?}"
1254 );
1255 }
1256
1257 const IDLE: Duration = Duration::from_secs(10);
1263
1264 async fn stream_after_one_event() -> (mpsc::Sender<BodyChunk>, EventStream) {
1266 let (tx, rx) = mpsc::channel(8);
1267 let mut stream = EventStream::new(rx)
1268 .with_jsonrpc_envelope(false)
1269 .with_idle_timeout(Some(IDLE));
1270 let event = make_status_event(TaskState::Working, false);
1271 tx.send(Ok(Bytes::from(bare_sse_frame(&event))))
1272 .await
1273 .unwrap();
1274 assert!(matches!(stream.next().await, Some(Ok(_))));
1275 (tx, stream)
1276 }
1277
1278 #[tokio::test(start_paused = true)]
1282 async fn idle_timeout_fires_after_the_first_event() {
1283 let (_tx, mut stream) = stream_after_one_event().await;
1284 let start = tokio::time::Instant::now();
1285 let result = stream.next().await;
1286 assert_eq!(start.elapsed(), IDLE, "fires exactly at the bound");
1287 match result {
1288 Some(Err(ClientError::Timeout(msg))) => {
1289 assert!(msg.contains("idle timeout"), "{msg}");
1290 assert!(msg.contains("10s"), "names the bound: {msg}");
1291 }
1292 other => panic!("expected idle Timeout, got {other:?}"),
1293 }
1294 assert!(stream.next().await.is_none(), "the stream has ended");
1295 }
1296
1297 #[tokio::test(start_paused = true)]
1301 async fn cancelled_polls_do_not_restart_the_idle_clock() {
1302 let (_tx, mut stream) = stream_after_one_event().await;
1303 for _ in 0..3 {
1304 let poll = tokio::time::timeout(IDLE / 4, stream.next()).await;
1305 assert!(poll.is_err(), "still inside the bound: {poll:?}");
1306 }
1307 let start = tokio::time::Instant::now();
1310 let result = tokio::time::timeout(IDLE, stream.next())
1311 .await
1312 .expect("the original deadline falls inside this poll");
1313 assert_eq!(start.elapsed(), IDLE / 4);
1314 assert!(matches!(result, Some(Err(ClientError::Timeout(_)))));
1315 }
1316
1317 #[tokio::test(start_paused = true)]
1320 async fn a_keep_alive_comment_restarts_the_idle_clock() {
1321 let (tx, mut stream) = stream_after_one_event().await;
1322 let quiet = IDLE * 3 / 4;
1323 tokio::time::advance(quiet).await;
1324 tx.send(Ok(Bytes::from_static(b": keep-alive\n\n")))
1325 .await
1326 .unwrap();
1327 let poll = tokio::time::timeout(quiet, stream.next()).await;
1329 assert!(poll.is_err(), "the comment renewed the bound: {poll:?}");
1330 let result = stream.next().await;
1332 assert!(matches!(result, Some(Err(ClientError::Timeout(_)))));
1333 }
1334
1335 #[tokio::test(start_paused = true)]
1338 async fn an_absent_or_unrepresentable_idle_bound_never_fires() {
1339 for bound in [None, Some(Duration::MAX)] {
1340 let (tx, rx) = mpsc::channel(8);
1341 let mut stream = EventStream::new(rx)
1342 .with_jsonrpc_envelope(false)
1343 .with_idle_timeout(bound);
1344 let event = make_status_event(TaskState::Working, false);
1345 tx.send(Ok(Bytes::from(bare_sse_frame(&event))))
1346 .await
1347 .unwrap();
1348 assert!(matches!(stream.next().await, Some(Ok(_))));
1349 let poll = tokio::time::timeout(IDLE * 1000, stream.next()).await;
1350 assert!(poll.is_err(), "{bound:?} must not fire: {poll:?}");
1351 }
1352 }
1353
1354 #[tokio::test(start_paused = true)]
1356 async fn streams_carry_the_default_idle_bound() {
1357 let (tx, rx) = mpsc::channel(8);
1358 let mut stream = EventStream::new(rx).with_jsonrpc_envelope(false);
1359 let event = make_status_event(TaskState::Working, false);
1360 tx.send(Ok(Bytes::from(bare_sse_frame(&event))))
1361 .await
1362 .unwrap();
1363 assert!(matches!(stream.next().await, Some(Ok(_))));
1364 let start = tokio::time::Instant::now();
1365 let result = stream.next().await;
1366 assert_eq!(
1367 start.elapsed(),
1368 crate::config::DEFAULT_STREAM_IDLE_TIMEOUT,
1369 "the default bound applies"
1370 );
1371 assert!(matches!(result, Some(Err(ClientError::Timeout(_)))));
1372 }
1373
1374 #[tokio::test(start_paused = true)]
1377 async fn idle_expiry_aborts_the_body_reader() {
1378 let (tx, rx) = mpsc::channel::<BodyChunk>(8);
1379 let reader = tokio::spawn(async move {
1380 let event = make_status_event(TaskState::Working, false);
1381 let _ = tx.send(Ok(Bytes::from(bare_sse_frame(&event)))).await;
1382 std::future::pending::<()>().await;
1383 });
1384 let mut stream = EventStream::with_abort_handle(rx, reader.abort_handle())
1385 .with_jsonrpc_envelope(false)
1386 .with_idle_timeout(Some(IDLE));
1387 assert!(matches!(stream.next().await, Some(Ok(_))));
1388 assert!(matches!(
1389 stream.next().await,
1390 Some(Err(ClientError::Timeout(_)))
1391 ));
1392 let joined = reader.await;
1393 assert!(
1394 joined.is_err_and(|e| e.is_cancelled()),
1395 "the reader must be aborted at expiry while the stream is still alive"
1396 );
1397 drop(stream);
1398 }
1399
1400 #[tokio::test]
1403 async fn stream_transport_error_from_channel() {
1404 let (tx, rx) = mpsc::channel(8);
1405 let mut stream = EventStream::new(rx);
1406
1407 tx.send(Err(ClientError::HttpClient("connection reset".into())))
1409 .await
1410 .unwrap();
1411
1412 let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
1413 .await
1414 .expect("timed out")
1415 .unwrap();
1416 match result {
1417 Err(ClientError::HttpClient(msg)) => {
1418 assert!(msg.contains("connection reset"));
1419 }
1420 other => panic!("expected HttpClient error, got {other:?}"),
1421 }
1422
1423 let end = tokio::time::timeout(TEST_TIMEOUT, stream.next())
1425 .await
1426 .expect("timed out");
1427 assert!(end.is_none(), "stream should end after transport error");
1428 }
1429
1430 #[tokio::test]
1431 async fn non_terminal_event_does_not_end_stream() {
1432 let (tx, rx) = mpsc::channel(8);
1433 let mut stream = EventStream::new(rx);
1434
1435 let working = make_status_event(TaskState::Working, false);
1437 let completed = make_status_event(TaskState::Completed, true);
1438 tx.send(Ok(Bytes::from(sse_frame(&working)))).await.unwrap();
1439 tx.send(Ok(Bytes::from(sse_frame(&completed))))
1440 .await
1441 .unwrap();
1442
1443 let first = tokio::time::timeout(TEST_TIMEOUT, stream.next())
1445 .await
1446 .expect("timed out on first event")
1447 .unwrap()
1448 .unwrap();
1449 assert!(
1450 matches!(first, StreamResponse::StatusUpdate(ref ev) if ev.status.state == TaskState::Working)
1451 );
1452
1453 let second = tokio::time::timeout(TEST_TIMEOUT, stream.next())
1455 .await
1456 .expect("timed out on second event")
1457 .unwrap()
1458 .unwrap();
1459 assert!(
1460 matches!(second, StreamResponse::StatusUpdate(ref ev) if ev.status.state == TaskState::Completed)
1461 );
1462
1463 let end = tokio::time::timeout(TEST_TIMEOUT, stream.next())
1465 .await
1466 .expect("timed out waiting for stream end");
1467 assert!(end.is_none());
1468 }
1469
1470 fn bare_sse_frame(event: &StreamResponse) -> String {
1474 let json = serde_json::to_string(event).unwrap();
1475 format!("data: {json}\n\n")
1476 }
1477
1478 #[tokio::test]
1479 async fn bare_stream_delivers_events() {
1480 let (tx, rx) = mpsc::channel(8);
1481 let mut stream = EventStream::new(rx).with_jsonrpc_envelope(false);
1482
1483 let event = make_status_event(TaskState::Working, false);
1484 tx.send(Ok(Bytes::from(bare_sse_frame(&event))))
1485 .await
1486 .unwrap();
1487 drop(tx);
1488
1489 let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
1490 .await
1491 .expect("timed out")
1492 .unwrap()
1493 .unwrap();
1494 assert!(
1495 matches!(result, StreamResponse::StatusUpdate(ref ev) if ev.status.state == TaskState::Working)
1496 );
1497 }
1498
1499 #[tokio::test]
1500 async fn bare_stream_ends_on_terminal() {
1501 let (tx, rx) = mpsc::channel(8);
1502 let mut stream = EventStream::new(rx).with_jsonrpc_envelope(false);
1503
1504 let event = make_status_event(TaskState::Completed, true);
1505 tx.send(Ok(Bytes::from(bare_sse_frame(&event))))
1506 .await
1507 .unwrap();
1508
1509 let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
1510 .await
1511 .expect("timed out")
1512 .unwrap()
1513 .unwrap();
1514 assert!(
1515 matches!(result, StreamResponse::StatusUpdate(ref ev) if ev.status.state == TaskState::Completed)
1516 );
1517
1518 let end = tokio::time::timeout(TEST_TIMEOUT, stream.next())
1519 .await
1520 .expect("timed out");
1521 assert!(end.is_none(), "bare stream should end after terminal event");
1522 }
1523
1524 #[tokio::test]
1525 async fn bare_stream_rejects_jsonrpc_envelope() {
1526 let (tx, rx) = mpsc::channel(8);
1527 let mut stream = EventStream::new(rx).with_jsonrpc_envelope(false);
1528
1529 let event = make_status_event(TaskState::Working, false);
1531 let envelope_frame = sse_frame(&event); tx.send(Ok(Bytes::from(envelope_frame))).await.unwrap();
1533 drop(tx);
1534
1535 let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
1536 .await
1537 .expect("timed out")
1538 .unwrap();
1539 assert!(
1540 result.is_err(),
1541 "bare stream should reject JSON-RPC envelope as invalid"
1542 );
1543 }
1544
1545 #[tokio::test]
1546 async fn envelope_stream_rejects_bare_response() {
1547 let (tx, rx) = mpsc::channel(8);
1548 let mut stream = EventStream::new(rx); let event = make_status_event(TaskState::Working, false);
1552 let bare_frame = bare_sse_frame(&event);
1553 tx.send(Ok(Bytes::from(bare_frame))).await.unwrap();
1554 drop(tx);
1555
1556 let result = tokio::time::timeout(TEST_TIMEOUT, stream.next())
1557 .await
1558 .expect("timed out")
1559 .unwrap();
1560 assert!(
1561 result.is_err(),
1562 "envelope stream should reject bare StreamResponse"
1563 );
1564 }
1565
1566 #[tokio::test]
1567 async fn bare_stream_multiple_events() {
1568 let (tx, rx) = mpsc::channel(8);
1569 let mut stream = EventStream::new(rx).with_jsonrpc_envelope(false);
1570
1571 let working = make_status_event(TaskState::Working, false);
1572 let completed = make_status_event(TaskState::Completed, true);
1573 tx.send(Ok(Bytes::from(bare_sse_frame(&working))))
1574 .await
1575 .unwrap();
1576 tx.send(Ok(Bytes::from(bare_sse_frame(&completed))))
1577 .await
1578 .unwrap();
1579
1580 let first = tokio::time::timeout(TEST_TIMEOUT, stream.next())
1581 .await
1582 .expect("timed out")
1583 .unwrap()
1584 .unwrap();
1585 assert!(
1586 matches!(first, StreamResponse::StatusUpdate(ref ev) if ev.status.state == TaskState::Working)
1587 );
1588
1589 let second = tokio::time::timeout(TEST_TIMEOUT, stream.next())
1590 .await
1591 .expect("timed out")
1592 .unwrap()
1593 .unwrap();
1594 assert!(
1595 matches!(second, StreamResponse::StatusUpdate(ref ev) if ev.status.state == TaskState::Completed)
1596 );
1597
1598 let end = tokio::time::timeout(TEST_TIMEOUT, stream.next())
1599 .await
1600 .expect("timed out");
1601 assert!(end.is_none());
1602 }
1603}