af-llm 0.6.0

use std::io::{Read, Write};
use std::net::TcpListener;
use std::time::Duration;

use af_llm::{ChatMessage, CompletionRequest, LlmClient, LlmConfig, LlmError};

#[tokio::test]
async fn premature_http_stream_end_without_done_fails_closed() {
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let address = listener.local_addr().unwrap();
    let server = std::thread::spawn(move || {
        let (mut socket, _) = listener.accept().unwrap();
        let mut request = [0; 4096];
        let _ = socket.read(&mut request).unwrap();
        let body = b"data: {\"id\":\"x\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"},\"finish_reason\":\"stop\"}]}\n\n";
        write!(
            socket,
            "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
            body.len()
        )
        .unwrap();
        socket.write_all(body).unwrap();
    });

    let mut config = LlmConfig::new(format!("http://{address}/v1"), "");
    config.timeout = Duration::from_secs(2);
    let client = LlmClient::new(config).unwrap();
    let request = CompletionRequest::new("model", vec![ChatMessage::user("hello")]).stream(true);
    let error = client
        .complete_stream_single_attempt(&request, |_, _| {})
        .await
        .unwrap_err();

    server.join().unwrap();
    assert!(matches!(error, LlmError::StreamProtocol(message) if message.contains("[DONE]")));
}

struct ImageSigner;
#[async_trait::async_trait]
impl af_llm::images::ImageResolver for ImageSigner {
    async fn resolve(
        &self,
        _: &af_context::RequestContext,
        _: &af_llm::InputImage,
    ) -> af_llm::Result<String> {
        Ok("https://assets.example/image?signature=private".into())
    }
}

#[tokio::test]
async fn images_reach_provider_as_content_parts_and_echoed_urls_are_redacted() {
    let listener = TcpListener::bind("127.0.0.1:0").unwrap();
    let address = listener.local_addr().unwrap();
    let server = std::thread::spawn(move || {
        let (mut socket, _) = listener.accept().unwrap();
        socket
            .set_read_timeout(Some(Duration::from_secs(3)))
            .unwrap();
        let mut request = Vec::new();
        let mut buffer = [0; 4096];
        let body = loop {
            let n = socket.read(&mut buffer).unwrap();
            assert!(n > 0 && request.len() < 16384);
            request.extend_from_slice(&buffer[..n]);
            if let Some(end) = request.windows(4).position(|w| w == b"\r\n\r\n") {
                let headers = String::from_utf8_lossy(&request[..end]).to_ascii_lowercase();
                let length: usize = headers
                    .lines()
                    .find_map(|line| line.strip_prefix("content-length:"))
                    .unwrap()
                    .trim()
                    .parse()
                    .unwrap();
                if request.len() >= end + 4 + length {
                    break serde_json::from_slice::<serde_json::Value>(
                        &request[end + 4..end + 4 + length],
                    )
                    .unwrap();
                }
            }
        };
        assert_eq!(
            body["messages"][0]["content"][1]["image_url"]["url"],
            "https://assets.example/image?signature=private"
        );
        assert!(body["messages"][0].get("images").is_none());
        assert!(body.get("context").is_none());
        let echo = "https://assets.example/image?signature=private";
        write!(
            socket,
            "HTTP/1.1 400 Bad Request\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
            echo.len(),
            echo
        )
        .unwrap();
    });
    let client = LlmClient::new(LlmConfig::new(format!("http://{address}/v1"), ""))
        .unwrap()
        .with_image_resolver(std::sync::Arc::new(ImageSigner));
    let mut message = ChatMessage::user("Describe");
    message.images.push(af_llm::InputImage {
        asset_id: "asset".parse().unwrap(),
        media_type: "image/png".into(),
    });
    let mut request = CompletionRequest::new("vision", vec![message]).stream(true);
    request.context = Some(af_context::RequestContext {
        tenant_id: "tenant".parse().unwrap(),
        subject_id: "owner".parse().unwrap(),
        request_id: "request".parse().unwrap(),
        locale: "en".into(),
        roles: Default::default(),
        entitlements: Default::default(),
    });
    let error = client
        .complete_stream_single_attempt(&request, |_, _| {})
        .await
        .unwrap_err();
    server.join().unwrap();
    assert!(error.to_string().contains("redacted"));
    assert!(!error.to_string().contains("signature"));
    assert!(!serde_json::to_string(&request)
        .unwrap()
        .contains("signature"));
}