hot_dev/
streaming.rs

1use std::pin::Pin;
2use std::sync::Arc;
3use std::time::Duration;
4
5use futures_core::Stream;
6use futures_util::StreamExt;
7use reqwest::header::ACCEPT;
8use reqwest::Method;
9use serde_json::Value;
10
11use crate::error::parse_api_error;
12use crate::sse::consume_sse_blocks;
13use crate::transport::{enc, Transport};
14use crate::{Error, JsonObject, StreamEvent};
15
16/// A stream of Hot run-stream events.
17pub type EventStream = Pin<Box<dyn Stream<Item = crate::Result<StreamEvent>> + Send>>;
18
19/// Convenience accessors for [`StreamEvent`] wire-format fields.
20pub trait StreamEventExt {
21    /// The event's "type" field, or "" when absent.
22    fn event_type(&self) -> &str;
23    /// The event's "run" object, when present.
24    fn run(&self) -> Option<&JsonObject>;
25    /// run.run_id, when present.
26    fn run_id(&self) -> Option<&str>;
27}
28
29impl StreamEventExt for StreamEvent {
30    fn event_type(&self) -> &str {
31        self.get("type").and_then(Value::as_str).unwrap_or("")
32    }
33
34    fn run(&self) -> Option<&JsonObject> {
35        self.get("run").and_then(Value::as_object)
36    }
37
38    fn run_id(&self) -> Option<&str> {
39        self.run()?.get("run_id")?.as_str()
40    }
41}
42
43/// Opens an SSE request and yields decoded events. The stream ends after
44/// yielding an error.
45pub(crate) fn iter_stream(
46    transport: Arc<Transport>,
47    method: Method,
48    path: String,
49    body: Option<Value>,
50    query: Vec<(String, String)>,
51) -> EventStream {
52    Box::pin(async_stream::stream! {
53        let mut builder = transport
54            .request_builder(method, &path)
55            .header(ACCEPT, "text/event-stream");
56        if !query.is_empty() {
57            builder = builder.query(&query);
58        }
59        if let Some(body) = &body {
60            builder = builder.json(body);
61        }
62
63        let response = match builder.send().await {
64            Ok(response) => response,
65            Err(error) => {
66                yield Err(Error::Http(error));
67                return;
68            }
69        };
70        let status = response.status();
71        if status.is_client_error() || status.is_server_error() {
72            let headers = response.headers().clone();
73            let text = response.text().await.unwrap_or_default();
74            yield Err(Error::Api(parse_api_error(status.as_u16(), &text, &headers)));
75            return;
76        }
77
78        // Bytes are decoded through a byte buffer so multi-byte UTF-8
79        // sequences split across chunks survive intact.
80        let mut pending: Vec<u8> = Vec::new();
81        let mut buffer = String::new();
82        let mut chunks = response.bytes_stream();
83        while let Some(chunk) = chunks.next().await {
84            let chunk = match chunk {
85                Ok(chunk) => chunk,
86                Err(error) => {
87                    yield Err(Error::Http(error));
88                    return;
89                }
90            };
91            pending.extend_from_slice(&chunk);
92            let valid_len = match std::str::from_utf8(&pending) {
93                Ok(_) => pending.len(),
94                Err(error) => error.valid_up_to(),
95            };
96            buffer.push_str(std::str::from_utf8(&pending[..valid_len]).unwrap_or(""));
97            pending.drain(..valid_len);
98
99            let (events, rest) = consume_sse_blocks(&buffer);
100            buffer = rest;
101            for event in events {
102                yield Ok(event);
103            }
104        }
105
106        buffer.push_str("\n\n");
107        let (events, _) = consume_sse_blocks(&buffer);
108        for event in events {
109            yield Ok(event);
110        }
111    })
112}
113
114pub(crate) async fn wait_for_run_result(
115    transport: Arc<Transport>,
116    stream_id: &str,
117    event_id: &str,
118    timeout: Duration,
119    mut on_chunk: Option<&mut (dyn FnMut(&str) + Send)>,
120) -> crate::Result<JsonObject> {
121    let deadline = tokio::time::Instant::now() + timeout;
122    let mut current_run_id: Option<String> = None;
123    for attempts in 0..=5 {
124        let mut source = iter_stream(
125            transport.clone(),
126            Method::GET,
127            format!("/streams/{}/subscribe", enc(stream_id)),
128            None,
129            Vec::new(),
130        );
131        loop {
132            let item = match tokio::time::timeout_at(deadline, source.next()).await {
133                Ok(item) => item,
134                Err(_) => return Err(Error::Timeout),
135            };
136            let Some(item) = item else { break };
137            match item {
138                Ok(event) => {
139                    if let Some(run) =
140                        handle_run_event(&event, event_id, &mut current_run_id, &mut on_chunk)?
141                    {
142                        return Ok(run);
143                    }
144                }
145                Err(_error) if attempts < 5 => break,
146                Err(error) => return Err(error),
147            }
148        }
149        if attempts == 5 {
150            return Err(Error::Protocol("stream ended before run completed".into()));
151        }
152    }
153    unreachable!()
154}
155
156pub(crate) async fn wait_for_call_result(
157    transport: Arc<Transport>,
158    body: JsonObject,
159    timeout: Duration,
160    mut on_chunk: Option<&mut (dyn FnMut(&str) + Send)>,
161) -> crate::Result<JsonObject> {
162    let deadline = tokio::time::Instant::now() + timeout;
163    let mut stream_id: Option<String> = None;
164    let mut event_id: Option<String> = None;
165    let mut current_run_id: Option<String> = None;
166
167    for attempts in 0..=5 {
168        let (method, path, request_body) = match stream_id.as_deref() {
169            None => (
170                Method::POST,
171                "/streams/subscribe-with-event".to_string(),
172                Some(Value::Object(body.clone())),
173            ),
174            Some(stream_id) => (
175                Method::GET,
176                format!("/streams/{}/subscribe", enc(stream_id)),
177                None,
178            ),
179        };
180        let mut source = iter_stream(transport.clone(), method, path, request_body, Vec::new());
181        loop {
182            let item = match tokio::time::timeout_at(deadline, source.next()).await {
183                Ok(item) => item,
184                Err(_) => return Err(Error::Timeout),
185            };
186            let Some(item) = item else { break };
187            let event = match item {
188                Ok(event) => event,
189                Err(_error) if stream_id.is_some() && attempts < 5 => break,
190                Err(error) => return Err(error),
191            };
192            if event.event_type() == "event:published" {
193                stream_id = event
194                    .get("stream_id")
195                    .and_then(Value::as_str)
196                    .map(str::to_string);
197                event_id = event
198                    .get("event_id")
199                    .and_then(Value::as_str)
200                    .map(str::to_string);
201            }
202            if let Some(event_id) = event_id.as_deref() {
203                if let Some(run) =
204                    handle_run_event(&event, event_id, &mut current_run_id, &mut on_chunk)?
205                {
206                    return Ok(run);
207                }
208            }
209        }
210        if stream_id.is_none() || event_id.is_none() {
211            return Err(Error::Protocol(
212                "stream ended before event:published was received".into(),
213            ));
214        }
215        if attempts == 5 {
216            return Err(Error::Protocol("stream ended before run completed".into()));
217        }
218    }
219    unreachable!()
220}
221
222pub(crate) async fn wait_for_task(
223    transport: Arc<Transport>,
224    task_id: &str,
225    timeout: Duration,
226    max_attempts: u32,
227) -> crate::Result<JsonObject> {
228    let deadline = tokio::time::Instant::now() + timeout;
229    for attempts in 0..=max_attempts {
230        let mut source = iter_stream(
231            transport.clone(),
232            Method::GET,
233            format!("/tasks/{}/subscribe", enc(task_id)),
234            None,
235            Vec::new(),
236        );
237        loop {
238            let item = match tokio::time::timeout_at(deadline, source.next()).await {
239                Ok(item) => item,
240                Err(_) => return Err(Error::TaskTimeout),
241            };
242            let Some(item) = item else { break };
243            let event = match item {
244                Ok(event) => event,
245                Err(_error) if attempts < max_attempts => break,
246                Err(error) => return Err(error),
247            };
248            if event.event_type() != "task:update" {
249                continue;
250            }
251            let Some(task) = event.get("task").and_then(Value::as_object) else {
252                continue;
253            };
254            match task.get("status").and_then(Value::as_str) {
255                Some("completed") => return Ok(task.clone()),
256                Some("failed" | "cancelled" | "timed_out") => {
257                    let status = task
258                        .get("status")
259                        .and_then(Value::as_str)
260                        .unwrap_or("failed");
261                    let message = result_message(task.get("result"))
262                        .map(str::to_string)
263                        .unwrap_or_else(|| format!("task {status}"));
264                    return Err(Error::TaskFailed {
265                        message,
266                        task: task.clone(),
267                    });
268                }
269                _ => {}
270            }
271        }
272        if attempts == max_attempts {
273            return Err(Error::Protocol(
274                "task subscription ended before the task completed".into(),
275            ));
276        }
277    }
278    unreachable!()
279}
280
281pub(crate) async fn wait_for_run(
282    transport: Arc<Transport>,
283    run_id: &str,
284    timeout: Duration,
285    max_attempts: u32,
286) -> crate::Result<JsonObject> {
287    let deadline = tokio::time::Instant::now() + timeout;
288    for attempts in 0..=max_attempts {
289        let mut source = iter_stream(
290            transport.clone(),
291            Method::GET,
292            format!("/runs/{}/subscribe", enc(run_id)),
293            None,
294            Vec::new(),
295        );
296        loop {
297            let item = match tokio::time::timeout_at(deadline, source.next()).await {
298                Ok(item) => item,
299                Err(_) => return Err(Error::Timeout),
300            };
301            let Some(item) = item else { break };
302            let event = match item {
303                Ok(event) => event,
304                Err(_error) if attempts < max_attempts => break,
305                Err(error) => return Err(error),
306            };
307            if event.event_type() != "run:update" {
308                continue;
309            }
310            let Some(run) = event.get("run").and_then(Value::as_object) else {
311                continue;
312            };
313            match run.get("status").and_then(Value::as_str) {
314                Some("succeeded") => return Ok(run.clone()),
315                Some(status @ ("failed" | "cancelled")) => {
316                    let message = run_result_message(run, &format!("run {status}"));
317                    return Err(Error::RunWaitFailed {
318                        message,
319                        run: run.clone(),
320                    });
321                }
322                _ => {}
323            }
324        }
325        if attempts == max_attempts {
326            return Err(Error::Protocol(
327                "run subscription ended before the run completed".into(),
328            ));
329        }
330    }
331    unreachable!()
332}
333
334// Keep the public Error::Api(ApiError) shape stable for the 1.x SDK. Boxing
335// that variant only to shrink this internal Result would be a breaking change.
336#[allow(clippy::result_large_err)]
337fn handle_run_event(
338    event: &StreamEvent,
339    event_id: &str,
340    current_run_id: &mut Option<String>,
341    on_chunk: &mut Option<&mut (dyn FnMut(&str) + Send)>,
342) -> crate::Result<Option<JsonObject>> {
343    match event.event_type() {
344        "run:start" => {
345            if event_id_of_run(event.run()) == Some(event_id) {
346                *current_run_id = event.run_id().map(str::to_string);
347            }
348        }
349        "stream:data" => {
350            let matches_run = current_run_id.as_deref().is_some_and(|current| {
351                event.get("run_id").and_then(Value::as_str) == Some(current)
352            });
353            if matches_run {
354                if let Some(on_chunk) = on_chunk.as_deref_mut() {
355                    if let Some(text) = event
356                        .get("payload")
357                        .and_then(Value::as_object)
358                        .and_then(|payload| payload.get("text"))
359                        .and_then(Value::as_str)
360                    {
361                        on_chunk(text);
362                    }
363                }
364            }
365        }
366        event_type @ ("run:stop" | "run:fail" | "run:cancel") => {
367            let Some(run) = event.run() else {
368                return Ok(None);
369            };
370            if event_id_of_run(Some(run)) != Some(event_id) {
371                return Ok(None);
372            }
373            match event_type {
374                "run:fail" => return Err(Error::RunFailed(run_result_message(run, "run failed"))),
375                "run:cancel" => {
376                    return Err(Error::RunFailed(run_result_message(run, "run cancelled")))
377                }
378                _ => return Ok(Some(run.clone())),
379            }
380        }
381        _ => {}
382    }
383    Ok(None)
384}
385
386pub(crate) fn event_id_of_run(run: Option<&JsonObject>) -> Option<&str> {
387    run?.get("event_id")?.as_str()
388}
389
390/// Digs a human-readable message out of a run result value.
391fn result_message(value: Option<&Value>) -> Option<&str> {
392    match value? {
393        Value::String(text) if !text.is_empty() => Some(text),
394        Value::Object(object) => ["$err", "$val", "error", "message", "msg", "reason", "err"]
395            .iter()
396            .find_map(|key| result_message(object.get(*key))),
397        _ => None,
398    }
399}
400
401pub(crate) fn run_result_message(run: &JsonObject, fallback: &str) -> String {
402    if let Some(message) = result_message(run.get("result")) {
403        return message.to_string();
404    }
405    match run.get("status").and_then(Value::as_str) {
406        Some(status) if !status.is_empty() => format!("{fallback} ({status})"),
407        _ => fallback.to_string(),
408    }
409}
410
411/// Unwraps a `{"$ok": ...}` run result.
412pub(crate) fn extract_run_result(result: Option<Value>) -> Value {
413    match result {
414        Some(Value::Object(mut object)) if object.contains_key("$ok") => {
415            object.remove("$ok").unwrap_or(Value::Null)
416        }
417        Some(other) => other,
418        None => Value::Null,
419    }
420}
hot_dev/streaming.rs

hot_dev/
streaming.rs
