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>(value: &T) -> Result<Bytes, serde_json::Error> {
71 SSE_FRAME_BUF.with(|cell| {
72 let mut buf = cell.borrow_mut();
73 buf.clear();
74 buf.extend_from_slice(b"event: message\ndata: ");
75 serde_json::to_writer(&mut *buf, value)?;
76 buf.extend_from_slice(b"\n\n");
77 Ok(Bytes::from(buf.clone()))
78 })
79}
80
81#[must_use]
83pub const fn write_keep_alive() -> Bytes {
84 Bytes::from_static(b": keep-alive\n\n")
85}
86
87#[derive(Debug)]
91pub struct SseBodyWriter {
92 tx: tokio::sync::mpsc::Sender<Result<Frame<Bytes>, Infallible>>,
93}
94
95#[allow(clippy::result_unit_err)]
100impl SseBodyWriter {
101 pub async fn send_event(&self, event_type: &str, data: &str) -> Result<(), ()> {
107 let frame = Frame::data(write_event(event_type, data));
108 self.tx.send(Ok(frame)).await.map_err(|_| ())
109 }
110
111 async fn send_raw_frame(&self, bytes: Bytes) -> Result<(), ()> {
120 let frame = Frame::data(bytes);
121 self.tx.send(Ok(frame)).await.map_err(|_| ())
122 }
123
124 pub async fn send_keep_alive(&self) -> Result<(), ()> {
130 let frame = Frame::data(write_keep_alive());
131 self.tx.send(Ok(frame)).await.map_err(|_| ())
132 }
133
134 pub fn close(self) {
136 drop(self);
137 }
138}
139
140struct ChannelBody {
146 rx: tokio::sync::mpsc::Receiver<Result<Frame<Bytes>, Infallible>>,
147}
148
149impl hyper::body::Body for ChannelBody {
150 type Data = Bytes;
151 type Error = Infallible;
152
153 fn poll_frame(
154 mut self: Pin<&mut Self>,
155 cx: &mut Context<'_>,
156 ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
157 self.rx.poll_recv(cx)
158 }
159}
160
161fn stream_error_payload(
171 err: &a2a_protocol_types::error::A2aError,
172 jsonrpc_envelope_id: Option<&JsonRpcId>,
173) -> Result<String, serde_json::Error> {
174 jsonrpc_envelope_id.map_or_else(
175 || serde_json::to_string(err),
176 |id| {
177 let mut jsonrpc_error = JsonRpcError::new(err.code.as_i32(), err.message.clone());
178 jsonrpc_error.data.clone_from(&err.data);
181 serde_json::to_string(&JsonRpcErrorResponse::new(id.clone(), jsonrpc_error))
182 },
183 )
184}
185
186#[must_use]
203#[allow(clippy::too_many_lines)]
204pub fn build_sse_response(
205 mut reader: InMemoryQueueReader,
206 keep_alive_interval: Option<Duration>,
207 channel_capacity: Option<usize>,
208 jsonrpc_envelope_id: Option<JsonRpcId>,
209) -> hyper::Response<http_body_util::combinators::BoxBody<Bytes, Infallible>> {
210 trace_info!("building SSE response stream");
211 let interval = keep_alive_interval.unwrap_or(DEFAULT_KEEP_ALIVE);
212 let cap = channel_capacity.unwrap_or(DEFAULT_SSE_CHANNEL_CAPACITY);
213 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Frame<Bytes>, Infallible>>(cap);
214
215 let body_writer = SseBodyWriter { tx };
216
217 tokio::spawn(async move {
218 tokio::task::yield_now().await;
226
227 let keep_alive_deadline = tokio::time::sleep(interval);
236 tokio::pin!(keep_alive_deadline);
237
238 loop {
239 tokio::select! {
240 biased;
241
242 event = reader.read() => {
243 match event {
244 Some(Ok(stream_response)) => {
245 let frame_bytes = if let Some(ref envelope_id) = jsonrpc_envelope_id {
250 let envelope = JsonRpcSuccessResponse {
253 jsonrpc: JsonRpcVersion,
254 id: envelope_id.clone(),
255 result: stream_response,
256 };
257 build_sse_message_frame(&envelope)
258 } else {
259 build_sse_message_frame(&stream_response)
261 };
262 let frame_bytes = match frame_bytes {
263 Ok(b) => b,
264 Err(e) => {
265 let err = a2a_protocol_types::error::A2aError::internal(
270 format!("event serialization failed: {e}"),
271 );
272 if let Ok(data) =
273 stream_error_payload(&err, jsonrpc_envelope_id.as_ref())
274 {
275 let _ = body_writer.send_event("error", &data).await;
276 }
277 break;
278 }
279 };
280 if body_writer.send_raw_frame(frame_bytes).await.is_err() {
281 break;
282 }
283 keep_alive_deadline.as_mut().reset(
285 tokio::time::Instant::now() + interval,
286 );
287 }
288 Some(Err(e)) => {
289 let Ok(data) =
300 stream_error_payload(&e, jsonrpc_envelope_id.as_ref())
301 else {
302 break;
303 };
304 let _ = body_writer.send_event("error", &data).await;
305 break;
306 }
307 None => break,
308 }
309 }
310 () = &mut keep_alive_deadline => {
311 if body_writer.send_keep_alive().await.is_err() {
312 break;
313 }
314 keep_alive_deadline.as_mut().reset(
315 tokio::time::Instant::now() + interval,
316 );
317 }
318 }
319 }
320
321 drop(body_writer);
322 });
323
324 let body = ChannelBody { rx };
325
326 hyper::Response::builder()
327 .status(200)
328 .header("content-type", "text/event-stream")
329 .header("cache-control", "no-cache")
330 .header("transfer-encoding", "chunked")
331 .body(body.boxed())
332 .unwrap_or_else(|_| {
333 hyper::Response::new(
334 http_body_util::Full::new(Bytes::from_static(b"SSE response build error")).boxed(),
335 )
336 })
337}
338
339#[cfg(test)]
340mod tests {
341 use super::*;
342
343 #[test]
346 fn write_event_single_line_data() {
347 let frame = write_event("message", r#"{"hello":"world"}"#);
348 let expected = "event: message\ndata: {\"hello\":\"world\"}\n\n";
349 assert_eq!(
350 frame,
351 Bytes::from(expected),
352 "single-line data should produce one data: line"
353 );
354 }
355
356 #[test]
357 fn write_event_multiline_data() {
358 let frame = write_event("error", "line1\nline2\nline3");
359 let expected = "event: error\ndata: line1\ndata: line2\ndata: line3\n\n";
360 assert_eq!(
361 frame,
362 Bytes::from(expected),
363 "multiline data should produce separate data: lines"
364 );
365 }
366
367 #[test]
368 fn write_event_empty_data() {
369 let frame = write_event("ping", "");
370 let expected = "event: ping\n\n";
372 assert_eq!(
373 frame,
374 Bytes::from(expected),
375 "empty data should produce no data: lines"
376 );
377 }
378
379 #[test]
380 fn write_event_empty_event_type() {
381 let frame = write_event("", "payload");
382 let expected = "event: \ndata: payload\n\n";
383 assert_eq!(
384 frame,
385 Bytes::from(expected),
386 "empty event type should still produce valid SSE frame"
387 );
388 }
389
390 #[test]
393 fn write_keep_alive_format() {
394 let frame = write_keep_alive();
395 assert_eq!(
396 frame,
397 Bytes::from_static(b": keep-alive\n\n"),
398 "keep-alive should be an SSE comment terminated by double newline"
399 );
400 }
401
402 #[tokio::test]
405 async fn sse_body_writer_send_event_delivers_frame() {
406 let (tx, mut rx) = tokio::sync::mpsc::channel::<Result<Frame<Bytes>, Infallible>>(8);
407 let writer = SseBodyWriter { tx };
408
409 writer
410 .send_event("message", "hello")
411 .await
412 .expect("send_event should succeed while receiver is alive");
413
414 let received = tokio::time::timeout(std::time::Duration::from_secs(5), rx.recv())
420 .await
421 .expect("send_event must deliver a frame; a timeout here means it returned Ok without sending")
422 .expect("channel should still be open");
423 let frame = received.expect("frame result should be Ok");
424 let data = frame.into_data().expect("frame should be a data frame");
425 assert_eq!(
426 data,
427 write_event("message", "hello"),
428 "received frame should match write_event output"
429 );
430 }
431
432 #[tokio::test]
433 async fn sse_body_writer_send_keep_alive_delivers_comment() {
434 let (tx, mut rx) = tokio::sync::mpsc::channel::<Result<Frame<Bytes>, Infallible>>(8);
435 let writer = SseBodyWriter { tx };
436
437 writer
438 .send_keep_alive()
439 .await
440 .expect("send_keep_alive should succeed while receiver is alive");
441
442 let received = tokio::time::timeout(std::time::Duration::from_secs(5), rx.recv())
444 .await
445 .expect("send_keep_alive must deliver a frame; a timeout here means it returned Ok without sending")
446 .expect("channel should still be open");
447 let frame = received.expect("frame result should be Ok");
448 let data = frame.into_data().expect("frame should be a data frame");
449 assert_eq!(
450 data,
451 write_keep_alive(),
452 "should receive keep-alive comment"
453 );
454 }
455
456 #[tokio::test]
457 async fn sse_body_writer_send_fails_after_receiver_dropped() {
458 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Frame<Bytes>, Infallible>>(8);
459 let writer = SseBodyWriter { tx };
460 drop(rx);
461
462 let result = writer.send_event("message", "data").await;
463 assert!(
464 result.is_err(),
465 "send_event should return Err after receiver is dropped"
466 );
467 }
468
469 #[tokio::test]
470 async fn sse_body_writer_keep_alive_fails_after_receiver_dropped() {
471 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Frame<Bytes>, Infallible>>(8);
472 let writer = SseBodyWriter { tx };
473 drop(rx);
474
475 let result = writer.send_keep_alive().await;
476 assert!(
477 result.is_err(),
478 "send_keep_alive should return Err after receiver is dropped"
479 );
480 }
481
482 #[tokio::test]
483 async fn sse_body_writer_close_drops_sender() {
484 let (tx, mut rx) = tokio::sync::mpsc::channel::<Result<Frame<Bytes>, Infallible>>(8);
485 let writer = SseBodyWriter { tx };
486
487 writer.close();
488
489 let result = rx.recv().await;
490 assert!(
491 result.is_none(),
492 "receiver should return None after writer is closed"
493 );
494 }
495
496 #[tokio::test]
499 async fn build_sse_response_has_correct_headers() {
500 let (_writer, reader) = crate::streaming::event_queue::new_in_memory_queue();
501
502 let response = build_sse_response(reader, None, None, Some(Some(serde_json::json!(1))));
503
504 assert_eq!(response.status(), 200, "status should be 200 OK");
505 assert_eq!(
506 response
507 .headers()
508 .get("content-type")
509 .map(hyper::http::HeaderValue::as_bytes),
510 Some(b"text/event-stream".as_slice()),
511 "Content-Type should be text/event-stream"
512 );
513 assert_eq!(
514 response
515 .headers()
516 .get("cache-control")
517 .map(hyper::http::HeaderValue::as_bytes),
518 Some(b"no-cache".as_slice()),
519 "Cache-Control should be no-cache"
520 );
521 assert_eq!(
522 response
523 .headers()
524 .get("transfer-encoding")
525 .map(hyper::http::HeaderValue::as_bytes),
526 Some(b"chunked".as_slice()),
527 "Transfer-Encoding should be chunked"
528 );
529 }
530
531 #[tokio::test]
532 async fn build_sse_response_with_custom_keep_alive_and_capacity() {
533 let (_writer, reader) = crate::streaming::event_queue::new_in_memory_queue();
535
536 let response = build_sse_response(
537 reader,
538 Some(Duration::from_secs(5)),
539 Some(16),
540 Some(Some(serde_json::json!(1))),
541 );
542
543 assert_eq!(response.status(), 200);
544 assert_eq!(
545 response
546 .headers()
547 .get("content-type")
548 .map(hyper::http::HeaderValue::as_bytes),
549 Some(b"text/event-stream".as_slice()),
550 );
551 }
552
553 #[tokio::test]
565 async fn build_sse_response_client_disconnect_stops_stream() {
566 use crate::streaming::event_queue::EventQueueWriter;
567 use a2a_protocol_types::events::{StreamResponse, TaskStatusUpdateEvent};
568 use a2a_protocol_types::task::{ContextId, TaskId, TaskState, TaskStatus};
569
570 let event = || {
571 StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
572 task_id: TaskId::new("t1"),
573 context_id: ContextId::new("c1"),
574 status: TaskStatus {
575 state: TaskState::Working,
576 message: None,
577 timestamp: None,
578 },
579 metadata: None,
580 })
581 };
582
583 let (writer, reader) = crate::streaming::event_queue::new_in_memory_queue();
584 let response = build_sse_response(reader, None, None, Some(Some(serde_json::json!(1))));
585
586 drop(response);
588 tokio::time::sleep(Duration::from_millis(50)).await;
589
590 writer
593 .write(event())
594 .await
595 .expect("the reader is still subscribed until the loop notices");
596
597 let mut observed_shutdown = false;
601 for _ in 0..200 {
602 tokio::time::sleep(Duration::from_millis(10)).await;
603 if writer.write(event()).await.is_err() {
604 observed_shutdown = true;
605 break;
606 }
607 }
608 assert!(
609 observed_shutdown,
610 "the writer loop never dropped its reader after the client disconnected"
611 );
612 }
613
614 #[tokio::test]
615 async fn build_sse_response_ends_on_reader_close() {
616 use http_body_util::BodyExt;
618
619 let (writer, reader) = crate::streaming::event_queue::new_in_memory_queue();
620
621 drop(writer);
623
624 let mut response = build_sse_response(reader, None, None, Some(Some(serde_json::json!(1))));
625
626 let frame = response.body_mut().frame().await;
628 if let Some(Ok(_)) = frame {
630 let next = response.body_mut().frame().await;
632 assert!(
633 next.is_none() || matches!(next, Some(Ok(_))),
634 "stream should eventually end"
635 );
636 }
637 }
638
639 #[tokio::test]
640 async fn build_sse_response_streams_error_event() {
641 use a2a_protocol_types::error::A2aError;
643 use http_body_util::BodyExt;
644
645 let (tx, rx) = tokio::sync::broadcast::channel(8);
648 let reader = crate::streaming::event_queue::InMemoryQueueReader::new(rx);
649
650 let err = A2aError::internal("something broke");
651 tx.send(Err(err)).expect("send should succeed");
652 drop(tx);
653
654 let mut response = build_sse_response(reader, None, None, Some(Some(serde_json::json!(1))));
655
656 let frame = response
657 .body_mut()
658 .frame()
659 .await
660 .expect("should have a frame")
661 .expect("frame should be Ok");
662 let data = frame.into_data().expect("should be a data frame");
663 let text = String::from_utf8_lossy(&data);
664
665 assert!(
666 text.starts_with("event: error\n"),
667 "error event frame should start with 'event: error\\n', got: {text}"
668 );
669 }
670
671 #[tokio::test]
672 async fn build_sse_response_streams_events() {
673 use crate::streaming::event_queue::EventQueueWriter;
674 use a2a_protocol_types::events::{StreamResponse, TaskStatusUpdateEvent};
675 use a2a_protocol_types::task::{ContextId, TaskId, TaskState, TaskStatus};
676 use http_body_util::BodyExt;
677
678 let (writer, reader) = crate::streaming::event_queue::new_in_memory_queue();
679
680 let event = StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
681 task_id: TaskId::new("t1"),
682 context_id: ContextId::new("c1"),
683 status: TaskStatus {
684 state: TaskState::Working,
685 message: None,
686 timestamp: None,
687 },
688 metadata: None,
689 });
690
691 writer.write(event).await.expect("write should succeed");
693 drop(writer);
694
695 let mut response = build_sse_response(reader, None, None, Some(Some(serde_json::json!(1))));
696
697 let frame = response
699 .body_mut()
700 .frame()
701 .await
702 .expect("should have a frame")
703 .expect("frame should be Ok");
704 let data = frame.into_data().expect("should be a data frame");
705 let text = String::from_utf8_lossy(&data);
706
707 assert!(
708 text.starts_with("event: message\n"),
709 "SSE frame should start with 'event: message\\n', got: {text}"
710 );
711 assert!(
712 text.contains("data: "),
713 "SSE frame should contain a data: line"
714 );
715 assert!(
717 text.contains("\"jsonrpc\""),
718 "data should contain JSON-RPC envelope"
719 );
720 assert!(
721 text.contains("\"result\""),
722 "data should contain result field"
723 );
724 let json_part = text
726 .lines()
727 .find_map(|l| l.strip_prefix("data: "))
728 .expect("frame must carry a data line");
729 let envelope: serde_json::Value =
730 serde_json::from_str(json_part).expect("data must be valid JSON");
731 assert_eq!(
732 envelope["id"],
733 serde_json::json!(1),
734 "SSE envelope must echo the request id, got: {envelope}"
735 );
736 }
737
738 #[tokio::test]
741 async fn build_sse_response_echoes_string_request_id() {
742 use crate::streaming::event_queue::EventQueueWriter;
743 use a2a_protocol_types::events::{StreamResponse, TaskStatusUpdateEvent};
744 use a2a_protocol_types::task::{ContextId, TaskId, TaskState, TaskStatus};
745 use http_body_util::BodyExt;
746
747 let (writer, reader) = crate::streaming::event_queue::new_in_memory_queue();
748 writer
749 .write(StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
750 task_id: TaskId::new("t1"),
751 context_id: ContextId::new("c1"),
752 status: TaskStatus {
753 state: TaskState::Working,
754 message: None,
755 timestamp: None,
756 },
757 metadata: None,
758 }))
759 .await
760 .expect("write should succeed");
761 drop(writer);
762
763 let mut response =
764 build_sse_response(reader, None, None, Some(Some(serde_json::json!("req-abc"))));
765 let frame = response
766 .body_mut()
767 .frame()
768 .await
769 .expect("should have a frame")
770 .expect("frame should be Ok");
771 let data = frame.into_data().expect("should be a data frame");
772 let text = String::from_utf8_lossy(&data);
773 let json_part = text
774 .lines()
775 .find_map(|l| l.strip_prefix("data: "))
776 .expect("frame must carry a data line");
777 let envelope: serde_json::Value =
778 serde_json::from_str(json_part).expect("data must be valid JSON");
779 assert_eq!(envelope["id"], serde_json::json!("req-abc"));
780 }
781}