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

cratestack_axum/rpc/
sse.rs
