path: root/crates/asist-core/src/http.rs

//! Minimal HTTP/1.1 client to talk to the local servers.
//!
//! The engines (llama-server, tts-server) listen on localhost over plain
//! HTTP, so neither TLS nor an async runtime is needed. In exchange for the
//! ~300 lines here we get the one thing no generic library makes easy and
//! this pipeline really needs: **aborting a response mid-download**. When the
//! user interrupts the assistant, reading already-generated audio must stop
//! and the socket must close that very instant; a client that only exposes
//! `read_to_end` cannot do that.

use std::io::{BufRead, BufReader, Read, Write};
use std::net::TcpStream;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;

use crate::error::{Error, Result};

/// Shared flag that aborts a read in progress.
///
/// It is checked between blocks, so the cut happens at most one `read` after
/// it is set, in the order of milliseconds with the block sizes the servers
/// use.
#[derive(Clone, Default)]
pub struct Cancel(Arc<AtomicBool>);

impl Cancel {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn cancel(&self) {
        self.0.store(true, Ordering::SeqCst);
    }

    pub fn is_cancelled(&self) -> bool {
        self.0.load(Ordering::SeqCst)
    }

    pub fn reset(&self) {
        self.0.store(false, Ordering::SeqCst);
    }
}

#[derive(Clone, Debug)]
pub struct HttpClient {
    /// `host:port` of the local server.
    pub authority: String,
    pub connect_timeout: Duration,
    /// Maximum time without receiving a single byte. It is not the total time:
    /// a long synthesis can take minutes as long as it keeps flowing.
    pub read_timeout: Duration,
}

impl HttpClient {
    pub fn new(authority: impl Into<String>) -> Self {
        Self {
            authority: authority.into(),
            connect_timeout: Duration::from_secs(5),
            read_timeout: Duration::from_secs(120),
        }
    }

    pub fn with_read_timeout(mut self, timeout: Duration) -> Self {
        self.read_timeout = timeout;
        self
    }

    fn connect(&self, path: &str) -> Result<TcpStream> {
        let addr = resolve(&self.authority, path)?;
        let stream = TcpStream::connect_timeout(&addr, self.connect_timeout).map_err(|e| {
            Error::Transport {
                url: format!("http://{}{}", self.authority, path),
                source: e,
            }
        })?;
        stream.set_read_timeout(Some(self.read_timeout))?;
        // Bodies are small and latency rules: no Nagle.
        let _ = stream.set_nodelay(true);
        Ok(stream)
    }

    fn send(&self, method: &str, path: &str, body: Option<&[u8]>) -> Result<Response> {
        let mut stream = self.connect(path)?;
        let mut head = format!(
            "{method} {path} HTTP/1.1\r\nHost: {}\r\nConnection: close\r\nAccept: */*\r\n",
            self.authority
        );
        if let Some(body) = body {
            head.push_str("Content-Type: application/json\r\n");
            head.push_str(&format!("Content-Length: {}\r\n", body.len()));
        }
        head.push_str("\r\n");

        let url = format!("http://{}{}", self.authority, path);
        let io = |e: std::io::Error| Error::Transport {
            url: url.clone(),
            source: e,
        };
        stream.write_all(head.as_bytes()).map_err(io)?;
        if let Some(body) = body {
            stream.write_all(body).map_err(io)?;
        }
        stream.flush().map_err(io)?;

        Response::read_head(BufReader::new(stream), url)
    }

    /// GET returning the whole, already decoded body.
    pub fn get(&self, path: &str) -> Result<Vec<u8>> {
        self.send("GET", path, None)?.into_body()
    }

    /// JSON POST returning the whole body.
    pub fn post_json(&self, path: &str, body: &serde_json::Value) -> Result<Vec<u8>> {
        let payload = serde_json::to_vec(body)?;
        self.send("POST", path, Some(&payload))?.into_body()
    }

    /// JSON POST whose response is consumed in chunks as it arrives.
    ///
    /// `on_chunk` gets each block as soon as it is available and returns
    /// `false` to stop; together with `cancel` they are the two ways barge-in
    /// stops a synthesis midway.
    pub fn post_json_streaming(
        &self,
        path: &str,
        body: &serde_json::Value,
        cancel: &Cancel,
        mut on_chunk: impl FnMut(&[u8]) -> bool,
    ) -> Result<()> {
        let payload = serde_json::to_vec(body)?;
        let mut response = self.send("POST", path, Some(&payload))?;
        response.error_for_status()?;
        response.stream(cancel, &mut on_chunk)
    }

    /// JSON POST that delivers the response line by line (llama-server SSE).
    pub fn post_json_lines(
        &self,
        path: &str,
        body: &serde_json::Value,
        cancel: &Cancel,
        mut on_line: impl FnMut(&str) -> bool,
    ) -> Result<()> {
        let mut pending = Vec::new();
        self.post_json_streaming(path, body, cancel, |chunk| {
            pending.extend_from_slice(chunk);
            while let Some(nl) = pending.iter().position(|b| *b == b'\n') {
                let line: Vec<u8> = pending.drain(..=nl).collect();
                let line = String::from_utf8_lossy(&line);
                if !on_line(line.trim_end_matches(['\r', '\n'])) {
                    return false;
                }
            }
            true
        })
    }

    /// Polls `/health` and returns `true` if the server is ready.
    pub fn healthy(&self, path: &str) -> bool {
        self.get(path).is_ok()
    }
}

fn resolve(authority: &str, path: &str) -> Result<std::net::SocketAddr> {
    use std::net::ToSocketAddrs;
    authority
        .to_socket_addrs()
        .map_err(|e| Error::Transport {
            url: format!("http://{authority}{path}"),
            source: e,
        })?
        .next()
        .ok_or_else(|| Error::Config(format!("could not resolve «{authority}»")))
}

/// Body encoding announced by the server.
enum Body {
    Chunked,
    Length(usize),
    /// No length and not chunked: the body ends when the socket closes.
    UntilClose,
}

pub struct Response {
    pub status: u16,
    url: String,
    reader: BufReader<TcpStream>,
    body: Body,
}

impl Response {
    fn read_head(mut reader: BufReader<TcpStream>, url: String) -> Result<Self> {
        let io = |e: std::io::Error| Error::Transport {
            url: url.clone(),
            source: e,
        };

        let mut line = String::new();
        reader.read_line(&mut line).map_err(io)?;
        let status = line
            .split_whitespace()
            .nth(1)
            .and_then(|s| s.parse::<u16>().ok())
            .ok_or_else(|| {
                Error::Config(format!("unreadable HTTP response from {url}: {line:?}"))
            })?;

        let mut body = Body::UntilClose;
        loop {
            let mut header = String::new();
            if reader.read_line(&mut header).map_err(io)? == 0 {
                break;
            }
            let header = header.trim_end();
            if header.is_empty() {
                break;
            }
            let Some((name, value)) = header.split_once(':') else {
                continue;
            };
            let (name, value) = (name.trim().to_ascii_lowercase(), value.trim());
            match name.as_str() {
                "transfer-encoding" if value.eq_ignore_ascii_case("chunked") => {
                    body = Body::Chunked;
                }
                "content-length" => {
                    if let Ok(n) = value.parse() {
                        body = Body::Length(n);
                    }
                }
                _ => {}
            }
        }

        Ok(Self {
            status,
            url,
            reader,
            body,
        })
    }

    fn error_for_status(&mut self) -> Result<()> {
        if