1use std::convert::Infallible;
12use std::pin::Pin;
13use std::task::{Context, Poll};
14use std::time::Duration;
15
16use bytes::Bytes;
17use http_body_util::BodyExt;
18use hyper::body::Frame;
19
20use a2a_protocol_types::jsonrpc::{
21 JsonRpcError, JsonRpcErrorResponse, JsonRpcId, JsonRpcSuccessResponse, JsonRpcVersion,
22};
23
24use crate::streaming::event_queue::{EventQueueReader, InMemoryQueueReader};
25
26pub(crate) const DEFAULT_KEEP_ALIVE: Duration = Duration::from_secs(30);
28
29pub(crate) const DEFAULT_SSE_CHANNEL_CAPACITY: usize = 64;
31
32#[must_use]
36pub fn write_event(event_type: &str, data: &str) -> Bytes {
37 let mut buf = String::with_capacity(event_type.len() + data.len() + 32);
38 buf.push_str("event: ");
39 buf.push_str(event_type);
40 buf.push('\n');
41 for line in data.lines() {
42 buf.push_str("data: ");
43 buf.push_str(line);
44 buf.push('\n');
45 }
46 buf.push('\n');
47 Bytes::from(buf)
48}
49
50std::thread_local! {
58 static SSE_FRAME_BUF: std::cell::RefCell<Vec<u8>> =
59 std::cell::RefCell::new(Vec::with_capacity(1024));
60}
61
62fn build_sse_message_frame<T: serde::Serialize>(
71 value: &T,
72 seq: Option<u64>,
73) -> Result<Bytes, serde_json::Error> {
74 SSE_FRAME_BUF.with(|cell| {
75 let mut buf = cell.borrow_mut();
76 buf.clear();
77 if let Some(seq) = seq {
87 buf.extend_from_slice(b"id: ");
88 push_decimal(&mut buf, seq);
89 buf.push(b'\n');
90 }
91 buf.extend_from_slice(b"event: message\ndata: ");
92 serde_json::to_writer(&mut *buf, value)?;
93 buf.extend_from_slice(b"\n\n");
94 Ok(Bytes::from(buf.clone()))
95 })
96}
97
98fn push_decimal(buf: &mut Vec<u8>, n: u64) {
106 let mut digits = [0u8; 20]; let mut i = digits.len();
108 let mut n = n;
109 loop {
110 i -= 1;
111 #[allow(clippy::cast_possible_truncation)]
113 {
114 digits[i] = b'0' + (n % 10) as u8;
115 }
116 n /= 10;
117 if n == 0 {
118 break;
119 }
120 }
121 buf.extend_from_slice(&digits[i..]);
122}
123
124#[must_use]
126pub const fn write_keep_alive() -> Bytes {
127 Bytes::from_static(b": keep-alive\n\n")
128}
129
130#[derive(Debug)]
134pub struct SseBodyWriter {
135 tx: tokio::sync::mpsc::Sender<Result<Frame<Bytes>, Infallible>>,
136}
137
138#[allow(clippy::result_unit_err)]
143impl SseBodyWriter {
144 pub async fn send_event(&self, event_type: &str, data: &str) -> Result<(), ()> {
150 let frame = Frame::data(write_event(event_type, data));
151 self.tx.send(Ok(frame)).await.map_err(|_| ())
152 }
153
154 async fn send_raw_frame(&self, bytes: Bytes) -> Result<(), ()> {
163 let frame = Frame::data(bytes);
164 self.tx.send(Ok(frame)).await.map_err(|_| ())
165 }
166
167 pub async fn send_keep_alive(&self) -> Result<(), ()> {
173 let frame = Frame::data(write_keep_alive());
174 self.tx.send(Ok(frame)).await.map_err(|_| ())
175 }
176
177 #[mutants::skip]
182 pub fn close(self) {
183 drop(self);
184 }
185}
186
187struct ChannelBody {
193 rx: tokio::sync::mpsc::Receiver<Result<Frame<Bytes>, Infallible>>,
194}
195
196impl hyper::body::Body for ChannelBody {
197 type Data = Bytes;
198 type Error = Infallible;
199
200 fn poll_frame(
201 mut self: Pin<&mut Self>,
202 cx: &mut Context<'_>,
203 ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
204 self.rx.poll_recv(cx)
205 }
206}
207
208fn stream_error_payload(
218 err: &a2a_protocol_types::error::A2aError,
219 jsonrpc_envelope_id: Option<&JsonRpcId>,
220) -> Result<String, serde_json::Error> {
221 jsonrpc_envelope_id.map_or_else(
222 || serde_json::to_string(err),
223 |id| {
224 let mut jsonrpc_error = JsonRpcError::new(err.code.as_i32(), err.message.clone());
225 jsonrpc_error.data.clone_from(&err.data);
228 serde_json::to_string(&JsonRpcErrorResponse::new(id.clone(), jsonrpc_error))
229 },
230 )
231}
232
233#[must_use]
250#[allow(clippy::too_many_lines)]
251pub fn build_sse_response(
252 mut reader: InMemoryQueueReader,
253 keep_alive_interval: Option<Duration>,
254 channel_capacity: Option<usize>,
255 jsonrpc_envelope_id: Option<JsonRpcId>,
256) -> hyper::Response<http_body_util::combinators::BoxBody<Bytes, Infallible>> {
257 trace_info!("building SSE response stream");
258 let interval = keep_alive_interval.unwrap_or(DEFAULT_KEEP_ALIVE);
259 let cap = channel_capacity.unwrap_or(DEFAULT_SSE_CHANNEL_CAPACITY);
260 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Frame<Bytes>, Infallible>>(cap);
261
262 let body_writer = SseBodyWriter { tx };
263
264 tokio::spawn(async move {
265 tokio::task::yield_now().await;
273
274 let keep_alive_deadline = tokio::time::sleep(interval);
283 tokio::pin!(keep_alive_deadline);
284
285 loop {
286 tokio::select! {
287 biased;
288
289 event = reader.read() => {
290 match event {
291 Some(Ok(queued)) => {
292 let seq = queued.seq;
293 let stream_response = queued.event;
294 let frame_bytes = if let Some(ref envelope_id) = jsonrpc_envelope_id {
299 let envelope = JsonRpcSuccessResponse {
302 jsonrpc: JsonRpcVersion,
303 id: envelope_id.clone(),
304 result: stream_response,
305 };
306 build_sse_message_frame(&envelope, seq)
307 } else {
308 build_sse_message_frame(&stream_response, seq)
310 };
311 let frame_bytes = match frame_bytes {
312 Ok(b) => b,
313 Err(e) => {
314 let err = a2a_protocol_types::error::A2aError::internal(
319 format!("event serialization failed: {e}"),
320 );
321 if let Ok(data) =
322 stream_error_payload(&err, jsonrpc_envelope_id.as_ref())
323 {
324 let _ = body_writer.send_event("error", &data).await;
325 }
326 break;
327 }
328 };
329 if body_writer.send_raw_frame(frame_bytes).await.is_err() {
330 break;
331 }
332 keep_alive_deadline.as_mut().reset(
334 tokio::time::Instant::now() + interval,
335 );
336 }
337 Some(Err(e)) => {
338 let Ok(data) =
349 stream_error_payload(&e, jsonrpc_envelope_id.as_ref())
350 else {
351 break;
352 };
353 let _ = body_writer.send_event("error", &data).await;
354 break;
355 }
356 None => break,
357 }
358 }
359 () = &mut keep_alive_deadline => {
360 if body_writer.send_keep_alive().await.is_err() {
361 break;
362 }
363 keep_alive_deadline.as_mut().reset(
364 tokio::time::Instant::now() + interval,
365 );
366 }
367 }
368 }
369
370 drop(body_writer);
371 });
372
373 let body = ChannelBody { rx };
374
375 hyper::Response::builder()
376 .status(200)
377 .header("content-type", "text/event-stream")
378 .header("cache-control", "no-cache")
379 .header("transfer-encoding", "chunked")
380 .body(body.boxed())
381 .unwrap_or_else(|_| {
382 hyper::Response::new(
383 http_body_util::Full::new(Bytes::from_static(b"SSE response build error")).boxed(),
384 )
385 })
386}
387
388#[cfg(test)]
389mod tests {
390 use super::*;
391
392 #[test]
395 fn write_event_single_line_data() {
396 let frame = write_event("message", r#"{"hello":"world"}"#);
397 let expected = "event: message\ndata: {\"hello\":\"world\"}\n\n";
398 assert_eq!(
399 frame,
400 Bytes::from(expected),
401 "single-line data should produce one data: line"
402 );
403 }
404
405 #[test]
406 fn write_event_multiline_data() {
407 let frame = write_event("error", "line1\nline2\nline3");
408 let expected = "event: error\ndata: line1\ndata: line2\ndata: line3\n\n";
409 assert_eq!(
410 frame,
411 Bytes::from(expected),
412 "multiline data should produce separate data: lines"
413 );
414 }
415
416 #[test]
417 fn write_event_empty_data() {
418 let frame = write_event("ping", "");
419 let expected = "event: ping\n\n";
421 assert_eq!(
422 frame,
423 Bytes::from(expected),
424 "empty data should produce no data: lines"
425 );
426 }
427
428 #[test]
429 fn write_event_empty_event_type() {
430 let frame = write_event("", "payload");
431 let expected = "event: \ndata: payload\n\n";
432 assert_eq!(
433 frame,
434 Bytes::from(expected),
435 "empty event type should still produce valid SSE frame"
436 );
437 }
438
439 #[test]
442 fn write_keep_alive_format() {
443 let frame = write_keep_alive();
444 assert_eq!(
445 frame,
446 Bytes::from_static(b": keep-alive\n\n"),
447 "keep-alive should be an SSE comment terminated by double newline"
448 );
449 }
450
451 #[tokio::test]
454 async fn sse_body_writer_send_event_delivers_frame() {
455 let (tx, mut rx) = tokio::sync::mpsc::channel::<Result<Frame<Bytes>, Infallible>>(8);
456 let writer = SseBodyWriter { tx };
457
458 writer
459 .send_event("message", "hello")
460 .await
461 .expect("send_event should succeed while receiver is alive");
462
463 let received = tokio::time::timeout(std::time::Duration::from_secs(5), rx.recv())
469 .await
470 .expect("send_event must deliver a frame; a timeout here means it returned Ok without sending")
471 .expect("channel should still be open");
472 let frame = received.expect("frame result should be Ok");
473 let data = frame.into_data().expect("frame should be a data frame");
474 assert_eq!(
475 data,
476 write_event("message", "hello"),
477 "received frame should match write_event output"
478 );
479 }
480
481 #[tokio::test]
482 async fn sse_body_writer_send_keep_alive_delivers_comment() {
483 let (tx, mut rx) = tokio::sync::mpsc::channel::<Result<Frame<Bytes>, Infallible>>(8);
484 let writer = SseBodyWriter { tx };
485
486 writer
487 .send_keep_alive()
488 .await
489 .expect("send_keep_alive should succeed while receiver is alive");
490
491 let received = tokio::time::timeout(std::time::Duration::from_secs(5), rx.recv())
493 .await
494 .expect("send_keep_alive must deliver a frame; a timeout here means it returned Ok without sending")
495 .expect("channel should still be open");
496 let frame = received.expect("frame result should be Ok");
497 let data = frame.into_data().expect("frame should be a data frame");
498 assert_eq!(
499 data,
500 write_keep_alive(),
501 "should receive keep-alive comment"
502 );
503 }
504
505 #[tokio::test]
506 async fn sse_body_writer_send_fails_after_receiver_dropped() {
507 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Frame<Bytes>, Infallible>>(8);
508 let writer = SseBodyWriter { tx };
509 drop(rx);
510
511 let result = writer.send_event("message", "data").await;
512 assert!(
513 result.is_err(),
514 "send_event should return Err after receiver is dropped"
515 );
516 }
517
518 #[tokio::test]
519 async fn sse_body_writer_keep_alive_fails_after_receiver_dropped() {
520 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Frame<Bytes>, Infallible>>(8);
521 let writer = SseBodyWriter { tx };
522 drop(rx);
523
524 let result = writer.send_keep_alive().await;
525 assert!(
526 result.is_err(),
527 "send_keep_alive should return Err after receiver is dropped"
528 );
529 }
530
531 #[tokio::test]
532 async fn sse_body_writer_close_drops_sender() {
533 let (tx, mut rx) = tokio::sync::mpsc::channel::<Result<Frame<Bytes>, Infallible>>(8);
534 let writer = SseBodyWriter { tx };
535
536 writer.close();
537
538 let result = rx.recv().await;
539 assert!(
540 result.is_none(),
541 "receiver should return None after writer is closed"
542 );
543 }
544
545 #[tokio::test]
548 async fn build_sse_response_has_correct_headers() {
549 let (_writer, reader) = crate::streaming::event_queue::new_in_memory_queue();
550
551 let response = build_sse_response(reader, None, None, Some(Some(serde_json::json!(1))));
552
553 assert_eq!(response.status(), 200, "status should be 200 OK");
554 assert_eq!(
555 response
556 .headers()
557 .get("content-type")
558 .map(hyper::http::HeaderValue::as_bytes),
559 Some(b"text/event-stream".as_slice()),
560 "Content-Type should be text/event-stream"
561 );
562 assert_eq!(
563 response
564 .headers()
565 .get("cache-control")
566 .map(hyper::http::HeaderValue::as_bytes),
567 Some(b"no-cache".as_slice()),
568 "Cache-Control should be no-cache"
569 );
570 assert_eq!(
571 response
572 .headers()
573 .get("transfer-encoding")
574 .map(hyper::http::HeaderValue::as_bytes),
575 Some(b"chunked".as_slice()),
576 "Transfer-Encoding should be chunked"
577 );
578 }
579
580 #[tokio::test]
581 async fn build_sse_response_with_custom_keep_alive_and_capacity() {
582 let (_writer, reader) = crate::streaming::event_queue::new_in_memory_queue();
584
585 let response = build_sse_response(
586 reader,
587 Some(Duration::from_secs(5)),
588 Some(16),
589 Some(Some(serde_json::json!(1))),
590 );
591
592 assert_eq!(response.status(), 200);
593 assert_eq!(
594 response
595 .headers()
596 .get("content-type")
597 .map(hyper::http::HeaderValue::as_bytes),
598 Some(b"text/event-stream".as_slice()),
599 );
600 }
601
602 #[tokio::test]
614 async fn build_sse_response_client_disconnect_stops_stream() {
615 use crate::streaming::event_queue::EventQueueWriter;
616 use a2a_protocol_types::events::{StreamResponse, TaskStatusUpdateEvent};
617 use a2a_protocol_types::task::{ContextId, TaskId, TaskState, TaskStatus};
618
619 let event = || {
620 StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
621 task_id: TaskId::new("t1"),
622 context_id: ContextId::new("c1"),
623 status: TaskStatus {
624 state: TaskState::Working,
625 message: None,
626 timestamp: None,
627 },
628 metadata: None,
629 })
630 };
631
632 let (writer, reader) = crate::streaming::event_queue::new_in_memory_queue();
633 let response = build_sse_response(reader, None, None, Some(Some(serde_json::json!(1))));
634
635 drop(response);
637 tokio::time::sleep(Duration::from_millis(50)).await;
638
639 writer
642 .write(event())
643 .await
644 .expect("the reader is still subscribed until the loop notices");
645
646 let mut observed_shutdown = false;
650 for _ in 0..200 {
651 tokio::time::sleep(Duration::from_millis(10)).await;
652 if writer.write(event()).await.is_err() {
653 observed_shutdown = true;
654 break;
655 }
656 }
657 assert!(
658 observed_shutdown,
659 "the writer loop never dropped its reader after the client disconnected"
660 );
661 }
662
663 #[tokio::test]
664 async fn build_sse_response_ends_on_reader_close() {
665 use http_body_util::BodyExt;
667
668 let (writer, reader) = crate::streaming::event_queue::new_in_memory_queue();
669
670 drop(writer);
672
673 let mut response = build_sse_response(reader, None, None, Some(Some(serde_json::json!(1))));
674
675 let frame = response.body_mut().frame().await;
677 if let Some(Ok(_)) = frame {
679 let next = response.body_mut().frame().await;
681 assert!(
682 next.is_none() || matches!(next, Some(Ok(_))),
683 "stream should eventually end"
684 );
685 }
686 }
687
688 #[tokio::test]
689 async fn build_sse_response_streams_error_event() {
690 use a2a_protocol_types::error::A2aError;
692 use http_body_util::BodyExt;
693
694 let (tx, rx) = tokio::sync::broadcast::channel(8);
697 let reader = crate::streaming::event_queue::InMemoryQueueReader::new(rx);
698
699 let err = A2aError::internal("something broke");
700 tx.send(Err(err)).expect("send should succeed");
701 drop(tx);
702
703 let mut response = build_sse_response(reader, None, None, Some(Some(serde_json::json!(1))));
704
705 let frame = response
706 .body_mut()
707 .frame()
708 .await
709 .expect("should have a frame")
710 .expect("frame should be Ok");
711 let data = frame.into_data().expect("should be a data frame");
712 let text = String::from_utf8_lossy(&data);
713
714 assert!(
715 text.starts_with("event: error\n"),
716 "error event frame should start with 'event: error\\n', got: {text}"
717 );
718 }
719
720 #[tokio::test]
721 async fn build_sse_response_streams_events() {
722 use crate::streaming::event_queue::EventQueueWriter;
723 use a2a_protocol_types::events::{StreamResponse, TaskStatusUpdateEvent};
724 use a2a_protocol_types::task::{ContextId, TaskId, TaskState, TaskStatus};
725 use http_body_util::BodyExt;
726
727 let (writer, reader) = crate::streaming::event_queue::new_in_memory_queue();
728
729 let event = StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
730 task_id: TaskId::new("t1"),
731 context_id: ContextId::new("c1"),
732 status: TaskStatus {
733 state: TaskState::Working,
734 message: None,
735 timestamp: None,
736 },
737 metadata: None,
738 });
739
740 writer.write(event).await.expect("write should succeed");
742 drop(writer);
743
744 let mut response = build_sse_response(reader, None, None, Some(Some(serde_json::json!(1))));
745
746 let frame = response
748 .body_mut()
749 .frame()
750 .await
751 .expect("should have a frame")
752 .expect("frame should be Ok");
753 let data = frame.into_data().expect("should be a data frame");
754 let text = String::from_utf8_lossy(&data);
755
756 assert!(
760 text.starts_with("id: 1\nevent: message\n"),
761 "SSE frame should open with its log position then the event line, got: {text}"
762 );
763 assert!(
764 text.contains("data: "),
765 "SSE frame should contain a data: line"
766 );
767 assert!(
769 text.contains("\"jsonrpc\""),
770 "data should contain JSON-RPC envelope"
771 );
772 assert!(
773 text.contains("\"result\""),
774 "data should contain result field"
775 );
776 let json_part = text
778 .lines()
779 .find_map(|l| l.strip_prefix("data: "))
780 .expect("frame must carry a data line");
781 let envelope: serde_json::Value =
782 serde_json::from_str(json_part).expect("data must be valid JSON");
783 assert_eq!(
784 envelope["id"],
785 serde_json::json!(1),
786 "SSE envelope must echo the request id, got: {envelope}"
787 );
788 }
789
790 #[tokio::test]
793 async fn build_sse_response_echoes_string_request_id() {
794 use crate::streaming::event_queue::EventQueueWriter;
795 use a2a_protocol_types::events::{StreamResponse, TaskStatusUpdateEvent};
796 use a2a_protocol_types::task::{ContextId, TaskId, TaskState, TaskStatus};
797 use http_body_util::BodyExt;
798
799 let (writer, reader) = crate::streaming::event_queue::new_in_memory_queue();
800 writer
801 .write(StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
802 task_id: TaskId::new("t1"),
803 context_id: ContextId::new("c1"),
804 status: TaskStatus {
805 state: TaskState::Working,
806 message: None,
807 timestamp: None,
808 },
809 metadata: None,
810 }))
811 .await
812 .expect("write should succeed");
813 drop(writer);
814
815 let mut response =
816 build_sse_response(reader, None, None, Some(Some(serde_json::json!("req-abc"))));
817 let frame = response
818 .body_mut()
819 .frame()
820 .await
821 .expect("should have a frame")
822 .expect("frame should be Ok");
823 let data = frame.into_data().expect("should be a data frame");
824 let text = String::from_utf8_lossy(&data);
825 let json_part = text
826 .lines()
827 .find_map(|l| l.strip_prefix("data: "))
828 .expect("frame must carry a data line");
829 let envelope: serde_json::Value =
830 serde_json::from_str(json_part).expect("data must be valid JSON");
831 assert_eq!(envelope["id"], serde_json::json!("req-abc"));
832 }
833
834 #[test]
841 fn a_positioned_event_carries_its_id_before_the_event_line() {
842 let frame =
843 build_sse_message_frame(&serde_json::json!({"a": 1}), Some(42)).expect("serialize");
844 assert_eq!(
845 String::from_utf8_lossy(&frame),
846 "id: 42\nevent: message\ndata: {\"a\":1}\n\n"
847 );
848 }
849
850 #[test]
855 fn an_unpositioned_event_carries_no_id_at_all() {
856 let frame = build_sse_message_frame(&serde_json::json!({"a": 1}), None).expect("serialize");
857 let text = String::from_utf8_lossy(&frame);
858 assert!(!text.contains("id:"), "no id line: {text}");
859 assert!(text.starts_with("event: message\n"), "{text}");
860 }
861
862 #[test]
866 fn the_reused_buffer_does_not_carry_an_id_into_the_next_frame() {
867 let _ = build_sse_message_frame(&serde_json::json!({"a": 1}), Some(7)).expect("first");
868 let second = build_sse_message_frame(&serde_json::json!({"a": 2}), None).expect("second");
869 assert!(
870 !String::from_utf8_lossy(&second).contains("id:"),
871 "the second frame must not inherit the first frame's id"
872 );
873 }
874
875 #[test]
876 fn positions_are_formatted_across_the_whole_u64_range() {
877 for (n, expected) in [
881 (0_u64, "0"),
882 (7, "7"),
883 (10, "10"),
884 (1_234_567_890, "1234567890"),
885 (u64::MAX, "18446744073709551615"),
886 ] {
887 let mut buf = Vec::new();
888 push_decimal(&mut buf, n);
889 assert_eq!(String::from_utf8_lossy(&buf), expected);
890 }
891 }
892}