audacity-sdk 0.2.0

//! Streaming output handle — wraps the SSE byte stream and exposes
//! `recv().await` (aws EventReceiver parity).

use futures_util::StreamExt;
use tokio::sync::mpsc;

use crate::mapping::{flush_metadata, process_chunk, StreamState};
use crate::sse::{SseParser, SseToken};
use crate::types::ConverseStreamOutput;
use crate::Error;

/// Handle returned by `converse_stream().send().await`.
///
/// Mirrors the aws-sdk-rust `EventReceiver` pattern:
/// ```rust,ignore
/// while let Some(event) = output.stream.recv().await? { … }
/// ```
pub struct ConverseStreamOutputHandle {
    pub stream: StreamReceiver,
}

/// The receiving end of the stream channel.
pub struct StreamReceiver {
    rx: mpsc::Receiver<Result<ConverseStreamOutput, Error>>,
}

impl StreamReceiver {
    /// Receive the next stream event.
    /// Returns `Ok(None)` when the stream has ended normally.
    /// Returns `Err(e)` on stream errors.
    pub async fn recv(&mut self) -> Result<Option<ConverseStreamOutput>, Error> {
        match self.rx.recv().await {
            Some(Ok(event)) => Ok(Some(event)),
            Some(Err(e)) => Err(e),
            None => Ok(None),
        }
    }
}

/// Spawn a background task that drives the SSE byte stream and sends parsed
/// events through the channel.
pub(crate) fn spawn_stream_driver(
    response: reqwest::Response,
    start_time: std::time::Instant,
) -> ConverseStreamOutputHandle {
    let (tx, rx) = mpsc::channel::<Result<ConverseStreamOutput, Error>>(64);

    tokio::spawn(async move {
        let mut byte_stream = response.bytes_stream();
        let mut parser = SseParser::new();
        let mut state = StreamState::new();
        state.start_time = start_time;

        let mut done = false; // saw the [DONE] sentinel
        let mut aborted = false; // already surfaced an error / receiver gone
        let mut eof = false; // byte stream ended

        'read: while !done && !eof {
            match byte_stream.next().await {
                Some(Ok(chunk)) => {
                    parser.push(&chunk);
                    if parser.overflowed() {
                        let _ = tx
                            .send(Err(Error::model_stream_error(
                                "SSE line exceeds the 32 MiB limit",
                            )))
                            .await;
                        aborted = true;
                        break 'read;
                    }
                }
                // Transport failure after the first SSE byte → spec §1/§4:
                // ModelStreamErrorException, never a raw I/O error, never a retry.
                Some(Err(e)) => {
                    let _ = tx
                        .send(Err(Error::model_stream_error(format!(
                            "stream read error: {e}"
                        ))))
                        .await;
                    aborted = true;
                    break 'read;
                }
                None => {
                    // Connection closed: flush a possible partial trailing line
                    // and run its tokens through the same handling below.
                    parser.flush();
                    eof = true;
                }
            }

            for token in parser.drain() {
                match token {
                    SseToken::Done => {
                        done = true;
                        break;
                    }
                    SseToken::Data(json_str) => match process_chunk(&mut state, &json_str) {
                        Ok(events) => {
                            for ev in events {
                                if tx.send(Ok(ev)).await.is_err() {
                                    aborted = true;
                                    break 'read;
                                }
                            }
                        }
                        Err(e) => {
                            let _ = tx.send(Err(e)).await;
                            aborted = true;
                            break 'read;
                        }
                    },
                }
            }
        }

        // Terminal rule: [DONE] seen → emit pending metadata as the final
        // event; connection dropped without [DONE] → ModelStreamError.
        if done {
            if let Some(meta) = flush_metadata(&state) {
                let _ = tx.send(Ok(meta)).await;
            }
        } else if !aborted {
            let _ = tx
                .send(Err(Error::model_stream_error(
                    "connection closed before [DONE]",
                )))
                .await;
        }
    });

    ConverseStreamOutputHandle {
        stream: StreamReceiver { rx },
    }
}