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

audacity_sdk/
stream.rs
