audacity_sdk/
stream.rs

1//! Streaming output handle — wraps the SSE byte stream and exposes
2//! `recv().await` (aws EventReceiver parity).
3
4use futures_util::StreamExt;
5use tokio::sync::mpsc;
6
7use crate::error::ErrorDetails;
8use crate::mapping::{flush_metadata, process_chunk, StreamState};
9use crate::sse::{SseParser, SseToken};
10use crate::types::ConverseStreamOutput;
11use crate::Error;
12
13/// Handle returned by `converse_stream().send().await`.
14///
15/// Mirrors the aws-sdk-rust `EventReceiver` pattern:
16/// ```rust,ignore
17/// while let Some(event) = output.stream.recv().await? { … }
18/// ```
19pub struct ConverseStreamOutputHandle {
20    pub stream: StreamReceiver,
21}
22
23/// The receiving end of the stream channel.
24pub struct StreamReceiver {
25    rx: mpsc::Receiver<Result<ConverseStreamOutput, Error>>,
26}
27
28impl StreamReceiver {
29    /// Receive the next stream event.
30    /// Returns `Ok(None)` when the stream has ended normally.
31    /// Returns `Err(e)` on stream errors.
32    pub async fn recv(&mut self) -> Result<Option<ConverseStreamOutput>, Error> {
33        match self.rx.recv().await {
34            Some(Ok(event)) => Ok(Some(event)),
35            Some(Err(e)) => Err(e),
36            None => Ok(None),
37        }
38    }
39}
40
41/// Spawn a background task that drives the SSE byte stream and sends parsed
42/// events through the channel.
43pub(crate) fn spawn_stream_driver(
44    response: reqwest::Response,
45    start_time: std::time::Instant,
46) -> ConverseStreamOutputHandle {
47    let (tx, rx) = mpsc::channel::<Result<ConverseStreamOutput, Error>>(64);
48
49    tokio::spawn(async move {
50        let mut byte_stream = response.bytes_stream();
51        let mut parser = SseParser::new();
52        let mut state = StreamState::new();
53        state.start_time = start_time;
54
55        let mut done = false; // saw the [DONE] sentinel
56        let mut aborted = false; // already surfaced an error / receiver gone
57        let mut eof = false; // byte stream ended
58
59        'read: while !done && !eof {
60            match byte_stream.next().await {
61                Some(Ok(chunk)) => parser.push(&chunk),
62                Some(Err(e)) => {
63                    let _ = tx
64                        .send(Err(Error::Sdk(format!("stream read error: {e}"))))
65                        .await;
66                    aborted = true;
67                    break 'read;
68                }
69                None => {
70                    // Connection closed: flush a possible partial trailing line
71                    // and run its tokens through the same handling below.
72                    parser.flush();
73                    eof = true;
74                }
75            }
76
77            for token in parser.drain() {
78                match token {
79                    SseToken::Done => {
80                        done = true;
81                        break;
82                    }
83                    SseToken::Data(json_str) => match process_chunk(&mut state, &json_str) {
84                        Ok(events) => {
85                            for ev in events {
86                                if tx.send(Ok(ev)).await.is_err() {
87                                    aborted = true;
88                                    break 'read;
89                                }
90                            }
91                        }
92                        Err(e) => {
93                            let _ = tx.send(Err(e)).await;
94                            aborted = true;
95                            break 'read;
96                        }
97                    },
98                }
99            }
100        }
101
102        // Terminal rule: [DONE] seen → emit pending metadata as the final
103        // event; connection dropped without [DONE] → ModelStreamError.
104        if done {
105            if let Some(meta) = flush_metadata(&state) {
106                let _ = tx.send(Ok(meta)).await;
107            }
108        } else if !aborted {
109            let _ = tx
110                .send(Err(Error::ModelStreamError(ErrorDetails {
111                    message: "connection closed before [DONE]".into(),
112                    status_code: 0,
113                    error_code: None,
114                    request_id: None,
115                    retry_after_seconds: None,
116                    raw_body: String::new(),
117                })))
118                .await;
119        }
120    });
121
122    ConverseStreamOutputHandle {
123        stream: StreamReceiver { rx },
124    }
125}
audacity_sdk/stream.rs

audacity_sdk/
stream.rs
