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
95impl SseBodyWriter {
96 pub async fn send_event(&self, event_type: &str, data: &str) -> Result<(), ()> {
102 let frame = Frame::data(write_event(event_type, data));
103 self.tx.send(Ok(frame)).await.map_err(|_| ())
104 }
105
106 async fn send_raw_frame(&self, bytes: Bytes) -> Result<(), ()> {
115 let frame = Frame::data(bytes);
116 self.tx.send(Ok(frame)).await.map_err(|_| ())
117 }
118
119 pub async fn send_keep_alive(&self) -> Result<(), ()> {
125 let frame = Frame::data(write_keep_alive());
126 self.tx.send(Ok(frame)).await.map_err(|_| ())
127 }
128
129 pub fn close(self) {
131 drop(self);
132 }
133}
134
135struct ChannelBody {
141 rx: tokio::sync::mpsc::Receiver<Result<Frame<Bytes>, Infallible>>,
142}
143
144impl hyper::body::Body for ChannelBody {
145 type Data = Bytes;
146 type Error = Infallible;
147
148 fn poll_frame(
149 mut self: Pin<&mut Self>,
150 cx: &mut Context<'_>,
151 ) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
152 self.rx.poll_recv(cx)
153 }
154}
155
156fn stream_error_payload(
166 err: &a2a_protocol_types::error::A2aError,
167 jsonrpc_envelope_id: Option<&JsonRpcId>,
168) -> Result<String, serde_json::Error> {
169 jsonrpc_envelope_id.map_or_else(
170 || serde_json::to_string(err),
171 |id| {
172 let mut jsonrpc_error = JsonRpcError::new(err.code.as_i32(), err.message.clone());
173 jsonrpc_error.data.clone_from(&err.data);
176 serde_json::to_string(&JsonRpcErrorResponse::new(id.clone(), jsonrpc_error))
177 },
178 )
179}
180
181#[must_use]
198#[allow(clippy::too_many_lines)]
199pub fn build_sse_response(
200 mut reader: InMemoryQueueReader,
201 keep_alive_interval: Option<Duration>,
202 channel_capacity: Option<usize>,
203 jsonrpc_envelope_id: Option<JsonRpcId>,
204) -> hyper::Response<http_body_util::combinators::BoxBody<Bytes, Infallible>> {
205 trace_info!("building SSE response stream");
206 let interval = keep_alive_interval.unwrap_or(DEFAULT_KEEP_ALIVE);
207 let cap = channel_capacity.unwrap_or(DEFAULT_SSE_CHANNEL_CAPACITY);
208 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Frame<Bytes>, Infallible>>(cap);
209
210 let body_writer = SseBodyWriter { tx };
211
212 tokio::spawn(async move {
213 tokio::task::yield_now().await;
221
222 let keep_alive_deadline = tokio::time::sleep(interval);
231 tokio::pin!(keep_alive_deadline);
232
233 loop {
234 tokio::select! {
235 biased;
236
237 event = reader.read() => {
238 match event {
239 Some(Ok(stream_response)) => {
240 let frame_bytes = if let Some(ref envelope_id) = jsonrpc_envelope_id {
245 let envelope = JsonRpcSuccessResponse {
248 jsonrpc: JsonRpcVersion,
249 id: envelope_id.clone(),
250 result: stream_response,
251 };
252 build_sse_message_frame(&envelope)
253 } else {
254 build_sse_message_frame(&stream_response)
256 };
257 let frame_bytes = match frame_bytes {
258 Ok(b) => b,
259 Err(e) => {
260 let err = a2a_protocol_types::error::A2aError::internal(
265 format!("event serialization failed: {e}"),
266 );
267 if let Ok(data) =
268 stream_error_payload(&err, jsonrpc_envelope_id.as_ref())
269 {
270 let _ = body_writer.send_event("error", &data).await;
271 }
272 break;
273 }
274 };
275 if body_writer.send_raw_frame(frame_bytes).await.is_err() {
276 break;
277 }
278 keep_alive_deadline.as_mut().reset(
280 tokio::time::Instant::now() + interval,
281 );
282 }
283 Some(Err(e)) => {
284 let Ok(data) =
295 stream_error_payload(&e, jsonrpc_envelope_id.as_ref())
296 else {
297 break;
298 };
299 let _ = body_writer.send_event("error", &data).await;
300 break;
301 }
302 None => break,
303 }
304 }
305 () = &mut keep_alive_deadline => {
306 if body_writer.send_keep_alive().await.is_err() {
307 break;
308 }
309 keep_alive_deadline.as_mut().reset(
310 tokio::time::Instant::now() + interval,
311 );
312 }
313 }
314 }
315
316 drop(body_writer);
317 });
318
319 let body = ChannelBody { rx };
320
321 hyper::Response::builder()
322 .status(200)
323 .header("content-type", "text/event-stream")
324 .header("cache-control", "no-cache")
325 .header("transfer-encoding", "chunked")
326 .body(body.boxed())
327 .unwrap_or_else(|_| {
328 hyper::Response::new(
329 http_body_util::Full::new(Bytes::from_static(b"SSE response build error")).boxed(),
330 )
331 })
332}
333
334#[cfg(test)]
335mod tests {
336 use super::*;
337
338 #[test]
341 fn write_event_single_line_data() {
342 let frame = write_event("message", r#"{"hello":"world"}"#);
343 let expected = "event: message\ndata: {\"hello\":\"world\"}\n\n";
344 assert_eq!(
345 frame,
346 Bytes::from(expected),
347 "single-line data should produce one data: line"
348 );
349 }
350
351 #[test]
352 fn write_event_multiline_data() {
353 let frame = write_event("error", "line1\nline2\nline3");
354 let expected = "event: error\ndata: line1\ndata: line2\ndata: line3\n\n";
355 assert_eq!(
356 frame,
357 Bytes::from(expected),
358 "multiline data should produce separate data: lines"
359 );
360 }
361
362 #[test]
363 fn write_event_empty_data() {
364 let frame = write_event("ping", "");
365 let expected = "event: ping\n\n";
367 assert_eq!(
368 frame,
369 Bytes::from(expected),
370 "empty data should produce no data: lines"
371 );
372 }
373
374 #[test]
375 fn write_event_empty_event_type() {
376 let frame = write_event("", "payload");
377 let expected = "event: \ndata: payload\n\n";
378 assert_eq!(
379 frame,
380 Bytes::from(expected),
381 "empty event type should still produce valid SSE frame"
382 );
383 }
384
385 #[test]
388 fn write_keep_alive_format() {
389 let frame = write_keep_alive();
390 assert_eq!(
391 frame,
392 Bytes::from_static(b": keep-alive\n\n"),
393 "keep-alive should be an SSE comment terminated by double newline"
394 );
395 }
396
397 #[tokio::test]
400 async fn sse_body_writer_send_event_delivers_frame() {
401 let (tx, mut rx) = tokio::sync::mpsc::channel::<Result<Frame<Bytes>, Infallible>>(8);
402 let writer = SseBodyWriter { tx };
403
404 writer
405 .send_event("message", "hello")
406 .await
407 .expect("send_event should succeed while receiver is alive");
408
409 let received = tokio::time::timeout(std::time::Duration::from_secs(5), rx.recv())
415 .await
416 .expect("send_event must deliver a frame; a timeout here means it returned Ok without sending")
417 .expect("channel should still be open");
418 let frame = received.expect("frame result should be Ok");
419 let data = frame.into_data().expect("frame should be a data frame");
420 assert_eq!(
421 data,
422 write_event("message", "hello"),
423 "received frame should match write_event output"
424 );
425 }
426
427 #[tokio::test]
428 async fn sse_body_writer_send_keep_alive_delivers_comment() {
429 let (tx, mut rx) = tokio::sync::mpsc::channel::<Result<Frame<Bytes>, Infallible>>(8);
430 let writer = SseBodyWriter { tx };
431
432 writer
433 .send_keep_alive()
434 .await
435 .expect("send_keep_alive should succeed while receiver is alive");
436
437 let received = tokio::time::timeout(std::time::Duration::from_secs(5), rx.recv())
439 .await
440 .expect("send_keep_alive must deliver a frame; a timeout here means it returned Ok without sending")
441 .expect("channel should still be open");
442 let frame = received.expect("frame result should be Ok");
443 let data = frame.into_data().expect("frame should be a data frame");
444 assert_eq!(
445 data,
446 write_keep_alive(),
447 "should receive keep-alive comment"
448 );
449 }
450
451 #[tokio::test]
452 async fn sse_body_writer_send_fails_after_receiver_dropped() {
453 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Frame<Bytes>, Infallible>>(8);
454 let writer = SseBodyWriter { tx };
455 drop(rx);
456
457 let result = writer.send_event("message", "data").await;
458 assert!(
459 result.is_err(),
460 "send_event should return Err after receiver is dropped"
461 );
462 }
463
464 #[tokio::test]
465 async fn sse_body_writer_keep_alive_fails_after_receiver_dropped() {
466 let (tx, rx) = tokio::sync::mpsc::channel::<Result<Frame<Bytes>, Infallible>>(8);
467 let writer = SseBodyWriter { tx };
468 drop(rx);
469
470 let result = writer.send_keep_alive().await;
471 assert!(
472 result.is_err(),
473 "send_keep_alive should return Err after receiver is dropped"
474 );
475 }
476
477 #[tokio::test]
478 async fn sse_body_writer_close_drops_sender() {
479 let (tx, mut rx) = tokio::sync::mpsc::channel::<Result<Frame<Bytes>, Infallible>>(8);
480 let writer = SseBodyWriter { tx };
481
482 writer.close();
483
484 let result = rx.recv().await;
485 assert!(
486 result.is_none(),
487 "receiver should return None after writer is closed"
488 );
489 }
490
491 #[tokio::test]
494 async fn build_sse_response_has_correct_headers() {
495 let (_writer, reader) = crate::streaming::event_queue::new_in_memory_queue();
496
497 let response = build_sse_response(reader, None, None, Some(Some(serde_json::json!(1))));
498
499 assert_eq!(response.status(), 200, "status should be 200 OK");
500 assert_eq!(
501 response
502 .headers()
503 .get("content-type")
504 .map(hyper::http::HeaderValue::as_bytes),
505 Some(b"text/event-stream".as_slice()),
506 "Content-Type should be text/event-stream"
507 );
508 assert_eq!(
509 response
510 .headers()
511 .get("cache-control")
512 .map(hyper::http::HeaderValue::as_bytes),
513 Some(b"no-cache".as_slice()),
514 "Cache-Control should be no-cache"
515 );
516 assert_eq!(
517 response
518 .headers()
519 .get("transfer-encoding")
520 .map(hyper::http::HeaderValue::as_bytes),
521 Some(b"chunked".as_slice()),
522 "Transfer-Encoding should be chunked"
523 );
524 }
525
526 #[tokio::test]
527 async fn build_sse_response_with_custom_keep_alive_and_capacity() {
528 let (_writer, reader) = crate::streaming::event_queue::new_in_memory_queue();
530
531 let response = build_sse_response(
532 reader,
533 Some(Duration::from_secs(5)),
534 Some(16),
535 Some(Some(serde_json::json!(1))),
536 );
537
538 assert_eq!(response.status(), 200);
539 assert_eq!(
540 response
541 .headers()
542 .get("content-type")
543 .map(hyper::http::HeaderValue::as_bytes),
544 Some(b"text/event-stream".as_slice()),
545 );
546 }
547
548 #[tokio::test]
549 async fn build_sse_response_client_disconnect_stops_stream() {
550 use crate::streaming::event_queue::EventQueueWriter;
552 use a2a_protocol_types::events::{StreamResponse, TaskStatusUpdateEvent};
553 use a2a_protocol_types::task::{ContextId, TaskId, TaskState, TaskStatus};
554
555 let (writer, reader) = crate::streaming::event_queue::new_in_memory_queue();
556
557 let response = build_sse_response(reader, None, None, Some(Some(serde_json::json!(1))));
558
559 drop(response);
561
562 tokio::time::sleep(Duration::from_millis(50)).await;
564
565 let event = StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
568 task_id: TaskId::new("t1"),
569 context_id: ContextId::new("c1"),
570 status: TaskStatus {
571 state: TaskState::Working,
572 message: None,
573 timestamp: None,
574 },
575 metadata: None,
576 });
577 let _ = writer.write(event).await;
579 drop(writer);
580 }
581
582 #[tokio::test]
583 async fn build_sse_response_ends_on_reader_close() {
584 use http_body_util::BodyExt;
586
587 let (writer, reader) = crate::streaming::event_queue::new_in_memory_queue();
588
589 drop(writer);
591
592 let mut response = build_sse_response(reader, None, None, Some(Some(serde_json::json!(1))));
593
594 let frame = response.body_mut().frame().await;
596 if let Some(Ok(_)) = frame {
598 let next = response.body_mut().frame().await;
600 assert!(
601 next.is_none() || matches!(next, Some(Ok(_))),
602 "stream should eventually end"
603 );
604 }
605 }
606
607 #[tokio::test]
608 async fn build_sse_response_streams_error_event() {
609 use a2a_protocol_types::error::A2aError;
611 use http_body_util::BodyExt;
612
613 let (tx, rx) = tokio::sync::broadcast::channel(8);
616 let reader = crate::streaming::event_queue::InMemoryQueueReader::new(rx);
617
618 let err = A2aError::internal("something broke");
619 tx.send(Err(err)).expect("send should succeed");
620 drop(tx);
621
622 let mut response = build_sse_response(reader, None, None, Some(Some(serde_json::json!(1))));
623
624 let frame = response
625 .body_mut()
626 .frame()
627 .await
628 .expect("should have a frame")
629 .expect("frame should be Ok");
630 let data = frame.into_data().expect("should be a data frame");
631 let text = String::from_utf8_lossy(&data);
632
633 assert!(
634 text.starts_with("event: error\n"),
635 "error event frame should start with 'event: error\\n', got: {text}"
636 );
637 }
638
639 #[tokio::test]
640 async fn build_sse_response_streams_events() {
641 use crate::streaming::event_queue::EventQueueWriter;
642 use a2a_protocol_types::events::{StreamResponse, TaskStatusUpdateEvent};
643 use a2a_protocol_types::task::{ContextId, TaskId, TaskState, TaskStatus};
644 use http_body_util::BodyExt;
645
646 let (writer, reader) = crate::streaming::event_queue::new_in_memory_queue();
647
648 let event = StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
649 task_id: TaskId::new("t1"),
650 context_id: ContextId::new("c1"),
651 status: TaskStatus {
652 state: TaskState::Working,
653 message: None,
654 timestamp: None,
655 },
656 metadata: None,
657 });
658
659 writer.write(event).await.expect("write should succeed");
661 drop(writer);
662
663 let mut response = build_sse_response(reader, None, None, Some(Some(serde_json::json!(1))));
664
665 let frame = response
667 .body_mut()
668 .frame()
669 .await
670 .expect("should have a frame")
671 .expect("frame should be Ok");
672 let data = frame.into_data().expect("should be a data frame");
673 let text = String::from_utf8_lossy(&data);
674
675 assert!(
676 text.starts_with("event: message\n"),
677 "SSE frame should start with 'event: message\\n', got: {text}"
678 );
679 assert!(
680 text.contains("data: "),
681 "SSE frame should contain a data: line"
682 );
683 assert!(
685 text.contains("\"jsonrpc\""),
686 "data should contain JSON-RPC envelope"
687 );
688 assert!(
689 text.contains("\"result\""),
690 "data should contain result field"
691 );
692 let json_part = text
694 .lines()
695 .find_map(|l| l.strip_prefix("data: "))
696 .expect("frame must carry a data line");
697 let envelope: serde_json::Value =
698 serde_json::from_str(json_part).expect("data must be valid JSON");
699 assert_eq!(
700 envelope["id"],
701 serde_json::json!(1),
702 "SSE envelope must echo the request id, got: {envelope}"
703 );
704 }
705
706 #[tokio::test]
709 async fn build_sse_response_echoes_string_request_id() {
710 use crate::streaming::event_queue::EventQueueWriter;
711 use a2a_protocol_types::events::{StreamResponse, TaskStatusUpdateEvent};
712 use a2a_protocol_types::task::{ContextId, TaskId, TaskState, TaskStatus};
713 use http_body_util::BodyExt;
714
715 let (writer, reader) = crate::streaming::event_queue::new_in_memory_queue();
716 writer
717 .write(StreamResponse::StatusUpdate(TaskStatusUpdateEvent {
718 task_id: TaskId::new("t1"),
719 context_id: ContextId::new("c1"),
720 status: TaskStatus {
721 state: TaskState::Working,
722 message: None,
723 timestamp: None,
724 },
725 metadata: None,
726 }))
727 .await
728 .expect("write should succeed");
729 drop(writer);
730
731 let mut response =
732 build_sse_response(reader, None, None, Some(Some(serde_json::json!("req-abc"))));
733 let frame = response
734 .body_mut()
735 .frame()
736 .await
737 .expect("should have a frame")
738 .expect("frame should be Ok");
739 let data = frame.into_data().expect("should be a data frame");
740 let text = String::from_utf8_lossy(&data);
741 let json_part = text
742 .lines()
743 .find_map(|l| l.strip_prefix("data: "))
744 .expect("frame must carry a data line");
745 let envelope: serde_json::Value =
746 serde_json::from_str(json_part).expect("data must be valid JSON");
747 assert_eq!(envelope["id"], serde_json::json!("req-abc"));
748 }
749}