quvyta-cli 0.1.0-alpha.5

//! A stand-in provider for tests: a real HTTP server on `127.0.0.1` that answers each request
//! with the next of the answers it was given and remembers what it was asked.
//!
//! It is a real socket rather than a stand-in transport so that the code that goes out over the
//! network is the code under test, and no test can reach a real provider: the only address a
//! test knows is this one.

use std::io::{BufRead, BufReader, Read, Write};
use std::net::TcpListener;
use std::sync::{Arc, Mutex, PoisonError};
use std::thread;

use serde_json::Value;

/// One request the server received.
#[derive(Debug, Clone)]
pub struct Received {
    /// The request line, such as `POST /v1/messages HTTP/1.1`.
    pub line: String,
    /// The headers, names in lower case.
    pub headers: Vec<(String, String)>,
    /// The body read as JSON.
    pub body: Value,
}

impl Received {
    /// The value of the header `name` (lower case), if it was sent.
    pub fn header(&self, name: &str) -> Option<&str> {
        self.headers.iter().find(|(key, _)| key == name).map(|(_, value)| value.as_str())
    }
}

/// A running stand-in provider.
pub struct FakeServer {
    /// The address to put in an endpoint.
    pub address: String,
    received: Arc<Mutex<Vec<Received>>>,
}

impl FakeServer {
    /// Starts a server that answers the requests in order with `answers`, each a status and a
    /// body; a request past the last answer is answered with 500.
    pub fn start(answers: Vec<(u16, String)>) -> Self {
        Self::start_holding(answers, std::time::Duration::ZERO)
    }

    /// Like [`start`](Self::start), but each answer's connection is held open for `hold` after the
    /// answer is written, the way a server that goes quiet in the middle of a reply holds it.
    pub fn start_holding(answers: Vec<(u16, String)>, hold: std::time::Duration) -> Self {
        let listener = TcpListener::bind("127.0.0.1:0").expect("a free port on the loopback");
        let address = format!("http://{}", listener.local_addr().expect("the port"));
        let received = Arc::new(Mutex::new(Vec::new()));
        let log = Arc::clone(&received);
        thread::spawn(move || {
            let mut answers = answers.into_iter();
            for stream in listener.incoming() {
                let Ok(mut stream) = stream else { return };
                let Some(request) = read_request(&mut stream) else { continue };
                log.lock().unwrap_or_else(PoisonError::into_inner).push(request);
                let (status, body) = answers.next().unwrap_or((500, "no more answers".to_owned()));
                let kind = if status == 200 { "text/event-stream" } else { "application/json" };
                let head = format!("HTTP/1.1 {status} Answer\r\ncontent-type: {kind}\r\nconnection: close\r\n\r\n");
                let _ = stream.write_all(head.as_bytes());
                let _ = stream.write_all(body.as_bytes());
                if !hold.is_zero() {
                    let _ = stream.flush();
                    thread::spawn(move || {
                        thread::sleep(hold);
                        drop(stream);
                    });
                }
            }
        });
        Self { address, received }
    }

    /// The requests received so far.
    pub fn received(&self) -> Vec<Received> {
        self.received.lock().unwrap_or_else(PoisonError::into_inner).clone()
    }
}

fn read_request(stream: &mut std::net::TcpStream) -> Option<Received> {
    let mut reader = BufReader::new(stream);
    let mut line = String::new();
    reader.read_line(&mut line).ok()?;
    let mut headers = Vec::new();
    loop {
        let mut header = String::new();
        reader.read_line(&mut header).ok()?;
        let header = header.trim_end();
        if header.is_empty() {
            break;
        }
        let (name, value) = header.split_once(':')?;
        headers.push((name.trim().to_ascii_lowercase(), value.trim().to_owned()));
    }
    let length: usize =
        headers.iter().find(|(name, _)| name == "content-length").and_then(|(_, value)| value.parse().ok())?;
    let mut body = vec![0; length];
    reader.read_exact(&mut body).ok()?;
    let body = serde_json::from_slice(&body).unwrap_or(Value::Null);
    Some(Received { line: line.trim_end().to_owned(), headers, body })
}

/// What a scripted reply says it used, the two counts a provider sends in `message_start` and
/// `message_delta`. A count left out is a server that did not say it, which qcli has to keep
/// going without rather than read as a zero.
#[derive(Debug, Clone, Copy)]
pub struct Usage {
    /// `message.usage.input_tokens`: the tokens the request read.
    pub read: Option<u64>,
    /// `usage.output_tokens`: the tokens the reply wrote.
    pub written: Option<u64>,
}

impl Usage {
    /// A reply that read `read` tokens and wrote `written` ones.
    pub fn new(read: u64, written: u64) -> Self {
        Self { read: Some(read), written: Some(written) }
    }

    /// A reply that counted nothing at all, the way a stream written by hand with [`sse`] can be.
    pub const SILENT: Self = Self { read: None, written: None };
}

/// The `message_start` that says what the request read, or nothing when it counted none: the
/// event only carries what the server said.
fn reads(read: Option<u64>) -> Option<Value> {
    read.map(|read| serde_json::json!({"type": "message_start", "message": {"usage": {"input_tokens": read}}}))
}

/// The `message_delta` that ends a reply: why it stopped, and what it wrote when it said. The
/// two travel in one event, the way a server sends them.
fn stops(reason: &str, written: Option<u64>) -> Value {
    let mut event = serde_json::json!({"type": "message_delta", "delta": {"stop_reason": reason}});
    if let Some(written) = written {
        event["usage"] = serde_json::json!({"output_tokens": written});
    }
    event
}

/// A streamed reply that says `text` and ends its turn, in pieces the way servers send them.
pub fn says(text: &str) -> String {
    says_using(text, Usage::new(10, 0))
}

/// A streamed reply like [`says`], but saying what it used, so a test can tell one turn's total
/// from another's.
pub fn says_using(text: &str, usage: Usage) -> String {
    let mut events = reads(usage.read).into_iter().collect::<Vec<_>>();
    events.push(
        serde_json::json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
    );
    for piece in text.split_inclusive(' ') {
        events.push(
            serde_json::json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": piece}}),
        );
    }
    events.push(serde_json::json!({"type": "content_block_stop", "index": 0}));
    events.push(stops("end_turn", usage.written));
    events.push(serde_json::json!({"type": "message_stop"}));
    sse(&events)
}

/// A streamed reply that calls tool `name` with `input` and waits for its result.
pub fn calls(id: &str, name: &str, input: &Value) -> String {
    calls_using(id, name, input, Usage::SILENT)
}

/// A streamed reply like [`calls`], but saying what it used, the way [`says_using`] does.
pub fn calls_using(id: &str, name: &str, input: &Value, usage: Usage) -> String {
    let mut events = reads(usage.read).into_iter().collect::<Vec<_>>();
    events.push(serde_json::json!({"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": id, "name": name, "input": {}}}));
    events.push(serde_json::json!({"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": input.to_string()}}));
    events.push(serde_json::json!({"type": "content_block_stop", "index": 0}));
    events.push(stops("tool_use", usage.written));
    events.push(serde_json::json!({"type": "message_stop"}));
    sse(&events)
}

/// `events` as a server-sent event stream.
pub fn sse(events: &[Value]) -> String {
    events.iter().map(|event| format!("event: {}\ndata: {event}\n\n", event["type"].as_str().unwrap_or(""))).collect()
}