cratestack_axum/rpc/
sse.rs

1//! SSE encoding for `@@subscribe` model-event streams (design doc
2//! §3.4a, cratestack#390). Counterpart to
3//! `crate::transport::stream_sequence`'s `application/cbor-seq` encoder
4//! for `@stream` procedures (§3.3) — same "encode item-by-item onto
5//! `axum::body::Body::from_stream`, never buffer the whole response"
6//! shape, but framed as `text/event-stream` instead of raw CBOR
7//! concatenation.
8//!
9//! Wire format: one `event: message` per item (`data:` is a JSON object
10//! `{"id": <u64>, "next": <item>}`, mirroring §2.3's `StreamItem` frame),
11//! and exactly one final `event: error` (`data:` is `{"id": <u64>,
12//! "err": RpcErrorBody}`) when the underlying stream ends — see
13//! `super::subscription_bridge`'s module doc for why "the stream ends"
14//! only ever means backpressure overflow here, never an ordinary client
15//! disconnect (which just drops this whole future instead, so this code
16//! never runs for that case). The generated `@@subscribe` endpoint now
17//! goes through [`super::sse_subscription`] instead, which takes the
18//! dispatcher's `SubscriptionStream` and sends its terminal error as is.
19//!
20//! Payload encoding is always JSON regardless of which `CratestackCodec` the
21//! server negotiates for its unary/batch RPC routes: SSE is a
22//! text-based wire format by construction, and JSON is already one of
23//! the two codecs this framework's RPC binding supports
24//! (`RPC_BINDING_CAPABILITIES`) — there's no reason to invent a
25//! base64-wrapped-CBOR convention nobody asked for when every
26//! off-the-shelf `EventSource` client and `curl` already expects JSON
27//! text bodies over SSE.
28
29use std::pin::Pin;
30
31use axum::body::{Body, Bytes};
32use axum::http::{HeaderMap, HeaderValue, StatusCode, header};
33use axum::response::Response;
34use cratestack_core::CratestackError;
35use cratestack_core::rpc::RpcErrorBody;
36use futures_util::Stream;
37use futures_util::stream::{self, StreamExt};
38use serde::Serialize;
39
40use crate::transport::StreamedResponseMarker;
41
42const SSE_CONTENT_TYPE: &str = "text/event-stream";
43
44/// `GET /rpc/subscribe/{op_id}` has no upgrade handshake to negotiate a
45/// binary subprotocol like WS does — the client states its intent via a
46/// plain `Accept` header, same as every other HTTP RPC binding. Reject
47/// anything that doesn't ask for SSE up front, before any
48/// `CratestackEventBus` subscription gets registered.
49pub fn validate_subscribe_accept_header(headers: &HeaderMap) -> Result<(), CratestackError> {
50    let Some(accept) = headers.get(header::ACCEPT) else {
51        return Err(CratestackError::NotAcceptable(format!(
52            "subscription endpoint requires Accept: {SSE_CONTENT_TYPE}"
53        )));
54    };
55    let accept = accept
56        .to_str()
57        .map_err(|error| CratestackError::BadRequest(format!("invalid Accept header: {error}")))?;
58    if accept
59        .split(',')
60        .map(str::trim)
61        .any(|value| value == SSE_CONTENT_TYPE || value == "*/*")
62    {
63        Ok(())
64    } else {
65        Err(CratestackError::NotAcceptable(format!(
66            "subscription endpoint requires Accept: {SSE_CONTENT_TYPE}, got {accept}"
67        )))
68    }
69}
70
71/// Encode `items` as a `text/event-stream` response. `items` ending
72/// (`None` from the underlying `Stream`) always means backpressure
73/// overflow closed the channel (see module docs) — the last byte chunk
74/// written is always the `Error{unavailable}` sentinel event.
75pub fn encode_model_event_sse_response<T, S>(items: S) -> Response
76where
77    T: Serialize + Send + 'static,
78    S: Stream<Item = T> + Send + 'static,
79{
80    sse_response(Body::from_stream(encode_sse_events(items)))
81}
82
83/// A `200 text/event-stream` response around an already-encoded body.
84pub(super) fn sse_response(body: Body) -> Response {
85    let mut response = Response::new(body);
86    *response.status_mut() = StatusCode::OK;
87    response.headers_mut().insert(
88        header::CONTENT_TYPE,
89        HeaderValue::from_static(SSE_CONTENT_TYPE),
90    );
91    response
92        .headers_mut()
93        .insert(header::CACHE_CONTROL, HeaderValue::from_static("no-cache"));
94    // See `crate::transport::StreamedResponseMarker` — bypasses any
95    // response-buffering middleware (e.g. `IdempotencyLayer`) an
96    // embedding app layers over the whole router, the same way the
97    // `@stream` cbor-seq encoder already does.
98    response.extensions_mut().insert(StreamedResponseMarker);
99    response
100}
101
102struct EventState<S> {
103    source: Pin<Box<S>>,
104    id: u64,
105    ended: bool,
106}
107
108fn encode_sse_events<T, S>(
109    items: S,
110) -> impl Stream<Item = Result<Bytes, std::convert::Infallible>> + Send
111where
112    T: Serialize + Send + 'static,
113    S: Stream<Item = T> + Send + 'static,
114{
115    let initial = EventState {
116        source: Box::pin(items),
117        id: 0,
118        ended: false,
119    };
120    stream::unfold(initial, |mut state| async move {
121        if state.ended {
122            return None;
123        }
124        match state.source.next().await {
125            Some(item) => {
126                state.id += 1;
127                let bytes = format_message_event(state.id, &item);
128                Some((Ok(bytes), state))
129            }
130            None => {
131                state.ended = true;
132                state.id += 1;
133                let bytes = format_error_event(state.id, &lagged_error());
134                Some((Ok(bytes), state))
135            }
136        }
137    })
138}
139
140fn lagged_error() -> RpcErrorBody {
141    RpcErrorBody::from_cratestack(&CratestackError::Unavailable(
142        "subscription lagged".to_owned(),
143    ))
144}
145
146#[derive(Serialize)]
147struct SseStreamItemPayload<'a, T> {
148    id: u64,
149    next: &'a T,
150}
151
152#[derive(Serialize)]
153struct SseErrorPayload<'a> {
154    id: u64,
155    err: &'a RpcErrorBody,
156}
157
158pub(super) fn format_message_event<T: Serialize>(id: u64, item: &T) -> Bytes {
159    let payload = SseStreamItemPayload { id, next: item };
160    let json =
161        serde_json::to_string(&payload).unwrap_or_else(|_| r#"{"id":0,"next":null}"#.to_owned());
162    Bytes::from(format!("event: message\ndata: {json}\n\n"))
163}
164
165pub(super) fn format_error_event(id: u64, error: &RpcErrorBody) -> Bytes {
166    let payload = SseErrorPayload { id, err: error };
167    let json = serde_json::to_string(&payload).unwrap_or_else(|_| {
168        r#"{"id":0,"err":{"code":"internal","message":"encode failure"}}"#.to_owned()
169    });
170    Bytes::from(format!("event: error\ndata: {json}\n\n"))
171}
172
173#[cfg(test)]
174mod tests;
cratestack_axum/rpc/sse.rs

cratestack_axum/rpc/
sse.rs
