saya-agent 0.3.1

use super::{framing::whitespace, ollama_chunks::Chunk, tool_assembly::ToolAssembly};
use crate::{CancellationToken, ProviderError, ProviderEvent, ProviderStream, TokenUsage};
use futures_util::{StreamExt, stream};
use reqwest::Response;
use std::{collections::VecDeque, time::Duration};

pub(super) fn parse(
    response: Response,
    cancellation: CancellationToken,
    idle: Duration,
) -> ProviderStream {
    Box::pin(stream::unfold(
        (
            response.bytes_stream(),
            State::default(),
            cancellation,
            idle,
        ),
        next,
    ))
}

async fn next<S>(
    mut value: (S, State, CancellationToken, Duration),
) -> Option<(
    Result<ProviderEvent, ProviderError>,
    (S, State, CancellationToken, Duration),
)>
where
    S: futures_util::Stream<Item = Result<bytes::Bytes, reqwest::Error>> + Unpin,
{
    loop {
        if value.2.is_cancelled() {
            value.1.done = true;
            return Some((Err(ProviderError::Cancelled), value));
        }
        if let Some(event) = value.1.pending.pop_front() {
            return Some((Ok(event), value));
        }
        if value.1.done {
            return None;
        }
        // Per-chunk idle budget instead of a total cap (see anthropic_stream).
        let item = tokio::select! {
            _ = value.2.cancelled() => return Some((Err(ProviderError::Cancelled), value)),
            item = tokio::time::timeout(value.3, value.0.next()) => match item {
                Ok(item) => item,
                Err(_) => {
                    value.1.done = true;
                    return Some((
                        Err(ProviderError::Request("provider stream stalled".into())),
                        value,
                    ));
                }
            }
        };
        let Some(chunk) = item else {
            if let Err(error) = value.1.finish() {
                value.1.done = true;
                return Some((Err(error), value));
            }
            continue;
        };
        let chunk = match chunk {
            Ok(chunk) => chunk,
            Err(_) => {
                value.1.done = true;
                return Some((
                    Err(ProviderError::Request("network request failed".into())),
                    value,
                ));
            }
        };
        if let Err(error) = value.1.push(&chunk) {
            value.1.done = true;
            return Some((Err(error), value));
        }
    }
}
#[derive(Default)]
struct State {
    bytes: Vec<u8>,
    usage: TokenUsage,
    pending: VecDeque<ProviderEvent>,
    tools: ToolAssembly,
    content: bool,
    done: bool,
}
impl State {
    fn finish(&mut self) -> Result<(), ProviderError> {
        if self.done {
            return if whitespace(&self.bytes) {
                self.bytes.clear();
                Ok(())
            } else {
                Err(ProviderError::InvalidResponse)
            };
        }
        let line = std::str::from_utf8(&self.bytes)
            .map_err(|_| ProviderError::InvalidResponse)?
            .to_owned();
        self.bytes.clear();
        self.push_line(line.trim_end_matches('\r'))?;
        if self.done {
            Ok(())
        } else {
            Err(ProviderError::InvalidResponse)
        }
    }
    fn push(&mut self, chunk: &[u8]) -> Result<(), ProviderError> {
        if self.bytes.len().saturating_add(chunk.len()) > crate::MAX_STREAM_BYTES {
            return Err(ProviderError::Request(
                "provider stream exceeded size limit".into(),
            ));
        }
        self.bytes.extend_from_slice(chunk);
        while let Some(end) = self.bytes.iter().position(|byte| *byte == b'\n') {
            let line = String::from_utf8(self.bytes[..end].to_vec())
                .map_err(|_| ProviderError::InvalidResponse)?;
            self.bytes.drain(..=end);
            self.push_line(line.trim_end_matches('\r'))?;
            if self.done {
                break;
            }
        }
        Ok(())
    }
    fn push_line(&mut self, line: &str) -> Result<(), ProviderError> {
        if line.trim().is_empty() {
            return Ok(());
        }
        let chunk: Chunk =
            serde_json::from_str(line).map_err(|_| ProviderError::InvalidResponse)?;
        if let Some(message) = chunk.message {
            if !message.content.is_empty() {
                self.content = true;
                self.pending
                    .push_back(ProviderEvent::TextDelta(message.content));
            }
            for (index, call) in message.tool_calls.into_iter().enumerate() {
                let arguments = match call.function.arguments {
                    serde_json::Value::String(value) => value,
                    value => value.to_string(),
                };
                self.tools.push(
                    index,
                    call.id.as_deref(),
                    Some(&call.function.name),
                    Some(&arguments),
                )?;
            }
        }
        if chunk.done {
            self.done = true;
            if let Some(input) = chunk.prompt_eval_count {
                self.usage.input_tokens = input;
            }
            if let Some(output) = chunk.eval_count {
                self.usage.output_tokens = output;
            }
            if self.usage != TokenUsage::default() {
                self.pending.push_back(ProviderEvent::Usage(self.usage));
            }
            self.complete()?;
        }
        Ok(())
    }
    fn complete(&mut self) -> Result<(), ProviderError> {
        if !whitespace(&self.bytes) || (!self.content && self.tools.is_empty()) {
            return Err(ProviderError::InvalidResponse);
        }
        let calls = std::mem::take(&mut self.tools).finish()?;
        if !calls.is_empty() {
            self.pending.push_back(ProviderEvent::ToolCalls(calls));
        }
        self.bytes.clear();
        self.pending.push_back(ProviderEvent::Done);
        Ok(())
    }
}