sie_sdk/client/
stream.rs

1//! Server-sent event streaming for the generation endpoints.
2//!
3//! Reconnect semantics mirror the buffered path with one addition: an error that arrives as
4//! the *first* SSE event is treated exactly like an HTTP 503, because nothing has been
5//! delivered yet and replaying is safe. Once a chunk has been yielded, any error is
6//! terminal: the caller has already consumed part of a generation that cannot be replayed.
7
8use std::pin::Pin;
9
10use bytes::Bytes;
11use futures_core::Stream;
12use futures_util::StreamExt;
13use reqwest::Method;
14use reqwest::header::HeaderMap;
15use serde::de::DeserializeOwned;
16use serde_json::Value;
17
18use crate::client::generate::{ChatRequest, GenerateRequest};
19use crate::client::{Client, StreamAttempt};
20use crate::error::{Error, Result, TransportErrorKind};
21use crate::http::{HttpResponse, PreparedRequest, headers};
22use crate::retry::{Decision, RequestOptions, RetryPolicy, RetryState};
23use crate::types::{ChatCompletionChunk, GenerateChunk};
24use crate::wire::sse::{self, LineDecoder};
25
26/// A stream of response chunks.
27///
28/// Boxed and pinned so it can be polled directly with [`futures_util::StreamExt`], without
29/// the caller having to pin it first.
30pub type ChunkStream<T> = Pin<Box<dyn Stream<Item = Result<T>> + Send>>;
31
32impl GenerateRequest {
33    /// Stream the generation token by token.
34    pub fn stream(self) -> Result<ChunkStream<GenerateChunk>> {
35        let body = self.body(true)?;
36        let request = self
37            .client
38            .sse_request(&self.path(), &body, &self.options)?;
39        Ok(Box::pin(sse_stream(
40            self.client,
41            request,
42            RetryPolicy::STREAM,
43            self.model,
44            self.options,
45        )))
46    }
47}
48
49impl ChatRequest {
50    /// Stream the completion token by token.
51    pub fn stream(self) -> Result<ChunkStream<ChatCompletionChunk>> {
52        let body = self.body(true)?;
53        let request = self
54            .client
55            .sse_request("/v1/chat/completions", &body, &self.options)?;
56        Ok(Box::pin(sse_stream(
57            self.client,
58            request,
59            RetryPolicy::STREAM,
60            self.model,
61            self.options,
62        )))
63    }
64}
65
66impl Client {
67    fn sse_request(
68        &self,
69        path: &str,
70        body: &Value,
71        options: &RequestOptions,
72    ) -> Result<PreparedRequest> {
73        let routing = self.routing(options.gpu.as_deref());
74        let encoded = serde_json::to_vec(body)
75            .map_err(|err| Error::invalid(format!("could not encode the request body: {err}")))?;
76        Ok(self
77            .request(Method::POST, path)?
78            .sse_headers()
79            .maybe_header(headers::MACHINE_PROFILE, routing.profile.as_deref())
80            .maybe_header(headers::POOL, routing.pool.as_deref())
81            .body(encoded))
82    }
83}
84
85/// A mid-stream capacity signal, rendered as the HTTP response it stands in for.
86///
87/// Reusing the response path means the stream's first-event errors go through exactly the
88/// same retry rules as a 503 received before the stream opened. The opening response's
89/// headers are carried over, so a `Retry-After` the gateway set on the connection is
90/// honoured rather than discarded.
91fn as_capacity_response(code: &str, message: &str, headers: &HeaderMap) -> HttpResponse {
92    let mut response_headers = headers.clone();
93    response_headers.insert(
94        reqwest::header::CONTENT_TYPE,
95        "application/json".parse().unwrap(),
96    );
97    HttpResponse {
98        status: 503,
99        headers: response_headers,
100        body: Bytes::from(
101            serde_json::json!({"error": {"code": code, "message": message}}).to_string(),
102        ),
103    }
104}
105
106/// What one decoded SSE event means for the stream.
107#[derive(Debug)]
108enum Event<T> {
109    Chunk(T),
110    /// The stream ended cleanly.
111    Done,
112    /// The server reported a failure before anything was delivered.
113    Capacity {
114        code: String,
115        message: String,
116    },
117}
118
119fn decode_event<T: DeserializeOwned>(payload: &str) -> Result<Event<T>> {
120    if payload == sse::DONE {
121        return Ok(Event::Done);
122    }
123    let value: Value = serde_json::from_str(payload)
124        .map_err(|err| Error::decode(format!("Malformed SSE chunk from server: {err}")))?;
125    if let Some((code, message)) = sse::chunk_error(&value) {
126        return Ok(Event::Capacity { code, message });
127    }
128    serde_json::from_value(value)
129        .map(Event::Chunk)
130        .map_err(|err| Error::decode(format!("Unexpected SSE chunk shape: {err}")))
131}
132
133fn sse_stream<T: DeserializeOwned + Send>(
134    client: Client,
135    request: PreparedRequest,
136    policy: RetryPolicy,
137    model: String,
138    options: RequestOptions,
139) -> impl Stream<Item = Result<T>> {
140    async_stream::try_stream! {
141        let mut state = RetryState::new(policy, &options, Some(&model));
142        let mut yielded = false;
143
144        loop {
145            let StreamAttempt { response, headers } = client.send_streaming(&request, &mut state).await?;
146            let mut body = response.bytes_stream();
147            let mut decoder = LineDecoder::default();
148            let mut reconnect_after = None;
149
150            'attempt: while let Some(chunk) = body.next().await {
151                // A transport failure part-way through a generation is never replayed.
152                let chunk = chunk.map_err(|error| {
153                    Error::connection(
154                        TransportErrorKind::MidFlight,
155                        format!("Connection lost during stream: {error}"),
156                        error,
157                    )
158                })?;
159
160                for line in decoder.push(&chunk) {
161                    let Some(payload) = sse::data_payload(&line) else { continue };
162                    match decode_event::<T>(payload)? {
163                        Event::Chunk(value) => {
164                            yielded = true;
165                            yield value;
166                        }
167                        Event::Done => return,
168                        Event::Capacity { code, message } => {
169                            let terminal = || Error::Server {
170                                message: message.clone(),
171                                code: Some(code.clone()),
172                                status: 503,
173                                request: None,
174                            };
175                            if yielded {
176                                Err(terminal())?;
177                            }
178                            match state.on_response(&as_capacity_response(&code, &message, &headers))? {
179                                Decision::Retry(delay) => {
180                                    reconnect_after = Some(delay);
181                                    break 'attempt;
182                                }
183                                // Unreachable: a 503 the policy does not retry raises above.
184                                Decision::Accept => Err(terminal())?,
185                            }
186                        }
187                    }
188                }
189            }
190
191            // A stream that ended without a trailing newline still has a final event.
192            if reconnect_after.is_none()
193                && let Some(line) = decoder.finish()
194                && let Some(payload) = sse::data_payload(&line)
195            {
196                match decode_event::<T>(payload)? {
197                    Event::Chunk(value) => yield value,
198                    Event::Done | Event::Capacity { .. } => return,
199                }
200            }
201
202            match reconnect_after {
203                // Dropping the response body here closes the connection, which the gateway
204                // sees as a client disconnect and stops generating for.
205                Some(delay) => tokio::time::sleep(delay).await,
206                // The stream ended without a `[DONE]` sentinel, which is still a clean end.
207                None => return,
208            }
209        }
210    }
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216    use crate::types::ChatMessage;
217
218    #[test]
219    fn decodes_payloads_sentinels_and_errors() {
220        let chunk: Event<ChatCompletionChunk> =
221            decode_event(r#"{"id": "c", "choices": [{"index": 0, "delta": {"content": "hi"}}]}"#)
222                .unwrap();
223        match chunk {
224            Event::Chunk(chunk) => assert_eq!(chunk.delta(), Some("hi")),
225            _ => panic!("expected a chunk"),
226        }
227
228        assert!(matches!(
229            decode_event::<ChatCompletionChunk>("[DONE]").unwrap(),
230            Event::Done
231        ));
232
233        let capacity: Event<ChatCompletionChunk> =
234            decode_event(r#"{"error": {"code": "RESOURCE_EXHAUSTED", "message": "oom"}}"#).unwrap();
235        match capacity {
236            Event::Capacity { code, message } => {
237                assert_eq!(code, "RESOURCE_EXHAUSTED");
238                assert_eq!(message, "oom");
239            }
240            _ => panic!("expected a capacity signal"),
241        }
242    }
243
244    #[test]
245    fn malformed_json_is_a_decode_error() {
246        let err = decode_event::<ChatCompletionChunk>("{not json").unwrap_err();
247        assert!(err.to_string().contains("Malformed SSE chunk"), "{err}");
248    }
249
250    #[test]
251    fn the_synthesized_capacity_response_keeps_the_connections_retry_hint() {
252        let mut headers = HeaderMap::new();
253        headers.insert(reqwest::header::RETRY_AFTER, "3".parse().unwrap());
254        let response = as_capacity_response("MODEL_LOADING", "loading", &headers);
255        assert_eq!(
256            crate::retry::backoff::retry_after(&response.headers),
257            Some(std::time::Duration::from_secs(3))
258        );
259    }
260
261    #[test]
262    fn the_synthesized_capacity_response_round_trips_through_the_error_reader() {
263        let response = as_capacity_response("MODEL_LOADING", "still loading", &HeaderMap::new());
264        assert_eq!(response.status, 503);
265        let envelope = crate::wire::parse_envelope(&response);
266        assert_eq!(envelope.code.as_deref(), Some("MODEL_LOADING"));
267        assert_eq!(envelope.message.as_deref(), Some("still loading"));
268    }
269
270    #[tokio::test]
271    async fn a_stream_against_an_unreachable_server_fails_rather_than_hanging() {
272        // Port 1 is reserved and never listening, so this exercises the connect path.
273        let client = Client::builder("http://127.0.0.1:1")
274            .timeout(std::time::Duration::from_millis(200))
275            .wait_for_capacity(false)
276            .build()
277            .unwrap();
278        let mut stream = client
279            .chat("m", [ChatMessage::user("hi")])
280            .stream()
281            .unwrap();
282        let first = stream
283            .next()
284            .await
285            .expect("the stream must yield a failure");
286        assert!(matches!(first, Err(Error::Connection { .. })), "{first:?}");
287    }
288}
sie_sdk/client/stream.rs

sie_sdk/client/
stream.rs
