aura-gui 0.2.2

//! Prueba del transporte completo contra un servidor local que responde SSE en trozos.

use aura_ui::ai::{AiError, Event, LanguageModel, Message, Provider, Request, StopReason};
use aura_ui::core::CancelToken;
use std::io::{Read, Write};
use std::net::TcpListener;
use std::thread;

fn request() -> Request {
    Request { model: "m".into(), system: None, messages: vec![Message::user("hola")], tools: vec![], max_tokens: 10 }
}

/// Atiende una conexión: lee la petición completa y responde con `status` y los trozos dados.
fn serve(status: &'static str, chunks: Vec<&'static str>) -> (String, thread::JoinHandle<String>) {
    let listener = TcpListener::bind("127.0.0.1:0").expect("puerto libre");
    let url = format!("http://{}", listener.local_addr().expect("dirección"));
    let handle = thread::spawn(move || {
        let (mut socket, _) = listener.accept().expect("conexión");
        let mut received = Vec::new();
        let mut buffer = [0u8; 4096];
        // Cabeceras más el cuerpo JSON, que cabe en un solo `read`.
        while !received.windows(4).any(|window| window == b"\r\n\r\n") {
            let read = socket.read(&mut buffer).expect("lectura");
            received.extend_from_slice(&buffer[..read]);
        }
        socket
            .write_all(
                format!("HTTP/1.1 {status}\r\ncontent-type: text/event-stream\r\nconnection: close\r\n\r\n").as_bytes(),
            )
            .expect("cabeceras");
        for chunk in chunks {
            socket.write_all(chunk.as_bytes()).expect("trozo");
            socket.flush().expect("flush");
            thread::sleep(std::time::Duration::from_millis(5));
        }
        String::from_utf8_lossy(&received).into_owned()
    });
    (url, handle)
}

#[test]
fn stream_openai_compatible_completo() {
    let (url, server) = serve(
        "200 OK",
        vec![
            "data: {\"choices\":[{\"delta\":{\"content\":\"Ho\"}}]}\n\ndata: {\"choi",
            "ces\":[{\"delta\":{\"content\":\"la\"},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n",
        ],
    );
    let provider = Provider { base_url: url, ..Provider::ollama() };
    let mut events = Vec::new();
    provider.stream(&request(), &CancelToken::default(), &mut |event| events.push(event)).expect("stream");
    let text: String = events
        .iter()
        .filter_map(|event| if let Event::Text(text) = event { Some(text.as_str()) } else { None })
        .collect();
    assert_eq!(text, "Hola");
    assert_eq!(events.last(), Some(&Event::Done(StopReason::EndTurn)));
    let received = server.join().expect("servidor");
    assert!(received.starts_with("POST /chat/completions"), "{received}");
}

#[test]
fn clave_viaja_en_la_cabecera_y_no_en_la_url() {
    let (url, server) = serve("200 OK", vec!["data: [DONE]\n\n"]);
    let provider = Provider { base_url: url, ..Provider::openai("sk-secreta".into()) };
    provider.stream(&request(), &CancelToken::default(), &mut |_| {}).expect("stream");
    let received = server.join().expect("servidor");
    let first_line = received.lines().next().unwrap_or_default();
    assert!(!first_line.contains("sk-secreta"));
    assert!(received.to_lowercase().contains("authorization: bearer sk-secreta"));
}

#[test]
fn error_http_devuelve_codigo_y_cuerpo() {
    let (url, _server) = serve("401 Unauthorized", vec!["{\"error\":\"clave inválida\"}"]);
    let provider = Provider { base_url: url, ..Provider::ollama() };
    let error = provider.stream(&request(), &CancelToken::default(), &mut |_| {}).unwrap_err();
    assert!(matches!(&error, AiError::Status { code: 401, body } if body.contains("clave inválida")), "{error}");
}

#[test]
fn cancelar_corta_el_stream() {
    let (url, _server) = serve("200 OK", vec!["data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n"; 50]);
    let provider = Provider { base_url: url, ..Provider::ollama() };
    let cancel = CancelToken::default();
    let mut seen = 0;
    let result = provider.stream(&request(), &cancel, &mut |_| {
        seen += 1;
        cancel.cancel();
    });
    assert!(matches!(result, Err(AiError::Cancelled)));
    assert!(seen < 50);
}

#[test]
fn cierre_antes_del_final_es_un_corte_y_no_una_respuesta_completa() {
    let (url, _server) = serve("200 OK", vec!["data: {\"choices\":[{\"delta\":{\"content\":\"Ho\"}}]}\n\n"]);
    let provider = Provider { base_url: url, ..Provider::ollama() };
    let mut events = Vec::new();
    let error = provider.stream(&request(), &CancelToken::default(), &mut |event| events.push(event)).unwrap_err();
    assert!(matches!(error, AiError::Transport(_)), "{error}");
    assert!(!events.iter().any(|event| matches!(event, Event::Done(_))), "no debe simular un final: {events:?}");
}

#[test]
fn anthropic_sin_message_stop_es_un_corte() {
    let (url, _server) = serve(
        "200 OK",
        vec![
            "event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"Ho\"}}\n\n",
        ],
    );
    let provider = Provider { base_url: url, ..Provider::anthropic("k".into()) };
    let error = provider.stream(&request(), &CancelToken::default(), &mut |_| {}).unwrap_err();
    assert!(matches!(error, AiError::Transport(_)), "{error}");
}

/// Atiende una conexión por cada respuesta dada `(estado, cabeceras extra, cuerpo)`, en orden.
fn serve_sequence(responses: Vec<(&'static str, &'static str, &'static str)>) -> (String, thread::JoinHandle<usize>) {
    let listener = TcpListener::bind("127.0.0.1:0").expect("puerto libre");
    let url = format!("http://{}", listener.local_addr().expect("dirección"));
    let handle = thread::spawn(move || {
        let mut served = 0;
        for (status, extra, body) in responses {
            let (mut socket, _) = listener.accept().expect("conexión");
            let mut received = Vec::new();
            let mut buffer = [0u8; 4096];
            // Cabeceras y cuerpo completos: cerrar con bytes sin leer provoca un reinicio de la conexión.
            loop {
                let read = socket.read(&mut buffer).expect("lectura");
                received.extend_from_slice(&buffer[..read]);
                let text = String::from_utf8_lossy(&received).to_lowercase();
                if let Some(end) = text.find("\r\n\r\n") {
                    let length: usize = text[..end]
                        .lines()
                        .find_map(|line| line.strip_prefix("content-length:"))
                        .map_or(0, |value| value.trim().parse().unwrap_or(0));
                    if received.len() >= end + 4 + length {
                        break;
                    }
                }
            }
            let head =
                format!("HTTP/1.1 {status}\r\ncontent-type: text/event-stream\r\n{extra}connection: close\r\n\r\n");
            socket.write_all(head.as_bytes()).expect("cabeceras");
            socket.write_all(body.as_bytes()).expect("cuerpo");
            served += 1;
        }
        served
    });
    (url, handle)
}

const DONE: &str = "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"},\"finish_reason\":\"stop\"}]}\n\n";

fn call(url: &str, policy: aura_ui::ai::http::RetryPolicy, cancel: &CancelToken) -> (Result<(), AiError>, Vec<Event>) {
    let mut events = Vec::new();
    let mut decoder = aura_ui::ai::openai::OpenAiDecoder::default();
    let result = aura_ui::ai::http::stream_sse_with(
        policy,
        &format!("{url}/chat/completions"),
        &[],
        &serde_json::json!({}),
        cancel,
        &mut decoder,
        &mut |event| events.push(event),
    );
    (result, events)
}

fn quick(retries: u32) -> aura_ui::ai::http::RetryPolicy {
    aura_ui::ai::http::RetryPolicy {
        retries,
        base_delay: std::time::Duration::from_millis(5),
        max_delay: std::time::Duration::from_millis(50),
    }
}

#[test]
fn reintenta_429_y_529_hasta_que_responde() {
    let (url, server) = serve_sequence(vec![
        ("429 Too Many Requests", "retry-after: 0\r\n", "{}"),
        ("529 Overloaded", "", "{}"),
        ("200 OK", "", DONE),
    ]);
    let (result, events) = call(&url, quick(3), &CancelToken::default());
    result.expect("debe recuperarse");
    assert_eq!(server.join().expect("servidor"), 3);
    assert!(events.contains(&Event::Text("ok".into())));
}

#[test]
fn al_agotar_los_reintentos_devuelve_el_ultimo_error() {
    let (url, server) = serve_sequence(vec![("503 Service Unavailable", "", "{}"); 3]);
    let (result, _) = call(&url, quick(2), &CancelToken::default());
    assert!(matches!(result, Err(AiError::Status { code: 503, .. })), "{result:?}");
    assert_eq!(server.join().expect("servidor"), 3, "1 intento + 2 reintentos");
}

#[test]
fn los_errores_del_cliente_no_se_reintentan() {
    let (url, server) = serve_sequence(vec![("401 Unauthorized", "", "{}")]);
    let (result, _) = call(&url, quick(3), &CancelToken::default());
    assert!(matches!(result, Err(AiError::Status { code: 401, .. })));
    assert_eq!(server.join().expect("servidor"), 1);
}

#[test]
fn cancelar_interrumpe_la_espera_entre_reintentos() {
    let (url, _server) = serve_sequence(vec![("503 Service Unavailable", "retry-after: 30\r\n", "{}"); 2]);
    let cancel = CancelToken::default();
    let remote = cancel.clone();
    thread::spawn(move || {
        thread::sleep(std::time::Duration::from_millis(100));
        remote.cancel();
    });
    let policy = aura_ui::ai::http::RetryPolicy { max_delay: std::time::Duration::from_secs(30), ..quick(3) };
    let started = std::time::Instant::now();
    let (result, _) = call(&url, policy, &cancel);
    assert!(matches!(result, Err(AiError::Cancelled)), "{result:?}");
    assert!(started.elapsed() < std::time::Duration::from_secs(5));
}

/// Un servidor que acepta y no contesta (una conexión reutilizada que quedó muerta): Detener libera al momento y
/// una respuesta que ya empezó no se corta por durar mucho, solo por un silencio largo.
#[test]
fn detener_no_espera_a_un_servidor_que_no_responde() {
    let listener = TcpListener::bind("127.0.0.1:0").expect("puerto libre");
    let url = format!("http://{}", listener.local_addr().expect("dirección"));
    let _server = thread::spawn(move || {
        let (_socket, _) = listener.accept().expect("conexión");
        thread::sleep(std::time::Duration::from_secs(10));
    });
    let provider = Provider { base_url: url, ..Provider::ollama() };
    let cancel = CancelToken::default();
    let stopper = cancel.clone();
    thread::spawn(move || {
        thread::sleep(std::time::Duration::from_millis(300));
        stopper.cancel();
    });
    let started = std::time::Instant::now();
    let error = provider.stream(&request(), &cancel, &mut |_| {}).unwrap_err();
    assert!(matches!(error, AiError::Cancelled), "{error}");
    assert!(started.elapsed() < std::time::Duration::from_secs(2), "{:?}", started.elapsed());
}