audacity-sdk 0.2.0

//! Integration tests against a local mock HTTP server.
//!
//! Tests spin up a real tokio TCP listener, serve canned HTTP responses,
//! and exercise the SDK end-to-end.

use std::sync::{
    atomic::{AtomicUsize, Ordering},
    Arc,
};

use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;

use audacity_sdk::{
    Client, Config, ContentBlock, ConversationRole, ConverseStreamOutput, Error, Message,
};

// ── helpers ───────────────────────────────────────────────────────────────────

/// Bind a random port and return (listener, base_url).
async fn bind() -> (TcpListener, String) {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let port = listener.local_addr().unwrap().port();
    (listener, format!("http://127.0.0.1:{port}"))
}

/// Build an HTTP/1.1 response string.
fn http_response(status: u16, status_text: &str, headers_extra: &str, body: &str) -> String {
    format!(
        "HTTP/1.1 {status} {status_text}\r\nContent-Type: application/json\r\nContent-Length: {len}\r\n{headers_extra}\r\n{body}",
        len = body.len()
    )
}

fn http_sse_response(body: &str) -> String {
    format!(
        "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nTransfer-Encoding: chunked\r\n\r\n{len:x}\r\n{body}\r\n0\r\n\r\n",
        len = body.len()
    )
}

/// Read one full HTTP request (headers + Content-Length body) from the socket.
async fn read_request(sock: &mut tokio::net::TcpStream) -> String {
    let mut buf: Vec<u8> = Vec::new();
    let mut tmp = [0u8; 4096];
    loop {
        let n = sock.read(&mut tmp).await.unwrap();
        if n == 0 {
            break;
        }
        buf.extend_from_slice(&tmp[..n]);
        if let Some(pos) = buf.windows(4).position(|w| w == b"\r\n\r\n") {
            let headers = String::from_utf8_lossy(&buf[..pos]).to_lowercase();
            let content_length = headers
                .lines()
                .find_map(|l| l.strip_prefix("content-length:"))
                .and_then(|v| v.trim().parse::<usize>().ok())
                .unwrap_or(0);
            if buf.len() >= pos + 4 + content_length {
                break;
            }
        }
    }
    String::from_utf8_lossy(&buf).to_string()
}

/// Serve a single connection with the given raw HTTP bytes.
async fn serve_once(listener: &TcpListener, response: String) {
    let (mut sock, _) = listener.accept().await.unwrap();
    read_request(&mut sock).await;
    sock.write_all(response.as_bytes()).await.unwrap();
}

fn client_for(base_url: &str) -> Client {
    Client::new(
        &Config::builder()
            .api_key("audacity_api_test")
            .base_url(base_url)
            .max_retries(0)
            .build()
            .unwrap(),
    )
    .unwrap()
}

fn user_message(text: &str) -> Message {
    Message::builder()
        .role(ConversationRole::User)
        .content(ContentBlock::Text(text.into()))
        .build()
        .unwrap()
}

// ── conformance §5.1 — converse happy path ────────────────────────────────────

#[tokio::test]
async fn converse_happy_path() {
    let (listener, base_url) = bind().await;

    let body = r#"{"choices":[{"message":{"role":"assistant","content":"Hello world"},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":3,"total_tokens":8}}"#;
    let response = http_response(200, "OK", "", body);

    let (result, _) = tokio::join!(
        async {
            let client = client_for(&base_url);
            client
                .converse()
                .model_id("gpt-5.4-mini")
                .messages(user_message("Hi"))
                .send()
                .await
        },
        serve_once(&listener, response),
    );

    let out = result.unwrap();
    let msg = out.output().unwrap().as_message().unwrap();
    assert_eq!(msg.content[0].as_text().unwrap(), "Hello world");
    assert_eq!(*out.stop_reason(), audacity_sdk::StopReason::EndTurn);
    assert_eq!(out.usage().input_tokens, 5);
    assert_eq!(out.usage().output_tokens, 3);
    assert_eq!(out.usage().total_tokens, 8);
    assert!(out.metrics().latency_ms < 5000);
}

// ── conformance §5.7 — defensive unwrap rule ─────────────────────────────────

#[tokio::test]
async fn defensive_unwrap_enveloped_body() {
    let (listener, base_url) = bind().await;

    let body = r#"{"data":{"choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{}}}"#;
    let response = http_response(200, "OK", "", body);

    let (result, _) = tokio::join!(
        async {
            client_for(&base_url)
                .converse()
                .model_id("m")
                .messages(user_message("Q"))
                .send()
                .await
        },
        serve_once(&listener, response),
    );

    assert!(result.unwrap().output().is_some());
}

// ── conformance §5.5 — error shape A ─────────────────────────────────────────

#[tokio::test]
async fn error_shape_a_access_denied() {
    let (listener, base_url) = bind().await;

    let body = r#"{"error":{"message":"Invalid API key","type":"authentication_error","code":"invalid_api_key"}}"#;
    let response = http_response(401, "Unauthorized", "", body);

    let (result, _) = tokio::join!(
        async {
            client_for(&base_url)
                .converse()
                .model_id("m")
                .messages(user_message("Q"))
                .send()
                .await
        },
        serve_once(&listener, response),
    );

    let err = result.unwrap_err();
    // §4: access-denied errors are not retryable — checked on the real enum.
    assert!(!err.is_retryable());
    match err {
        Error::AccessDenied(d) => {
            assert_eq!(d.error_code.as_deref(), Some("invalid_api_key"));
        }
        other => panic!("expected AccessDenied, got {other:?}"),
    }
}

// ── conformance §5.5 — error shape B ─────────────────────────────────────────

#[tokio::test]
async fn error_shape_b_model_not_allowed() {
    let (listener, base_url) = bind().await;

    let body = r#"{"success":false,"error":{"code":"MODEL_NOT_ALLOWED","message":"Not allowed","request_id":"req-abc","details":{}}}"#;
    let response = http_response(403, "Forbidden", "", body);

    let (result, _) = tokio::join!(
        async {
            client_for(&base_url)
                .converse()
                .model_id("m")
                .messages(user_message("Q"))
                .send()
                .await
        },
        serve_once(&listener, response),
    );

    assert!(matches!(result, Err(Error::AccessDenied(_))));
    if let Err(Error::AccessDenied(d)) = result {
        assert_eq!(d.request_id.as_deref(), Some("req-abc"));
    }
}

// ── conformance §5.6 — 429 with Retry-After retries then succeeds ─────────────

#[tokio::test]
async fn retry_after_429_then_success() {
    converse_with_429_then_ok(0).await;
}

/// Serve a 429 (with the given `Retry-After`) then a 200, drive one `converse`
/// with one retry allowed, assert it succeeds, and return the elapsed time.
async fn converse_with_429_then_ok(retry_after_secs: u64) -> std::time::Duration {
    let (listener, base_url) = bind().await;

    let throttle_body = r#"{"error":{"message":"rate limited","code":"rate_limit_exceeded"}}"#;
    let throttle_resp = http_response(
        429,
        "Too Many Requests",
        &format!("Retry-After: {retry_after_secs}\r\n"),
        throttle_body,
    );
    let ok_body = r#"{"choices":[{"message":{"content":"ok"},"finish_reason":"stop"}],"usage":{}}"#;
    let ok_resp = http_response(200, "OK", "", ok_body);

    let client = Client::new(
        &Config::builder()
            .api_key("audacity_api_test")
            .base_url(&base_url)
            .max_retries(1) // 2 attempts total
            .build()
            .unwrap(),
    )
    .unwrap();

    let serve_two = async {
        // first request → 429
        serve_once(&listener, throttle_resp).await;
        // second request → 200
        serve_once(&listener, ok_resp).await;
    };

    let started = std::time::Instant::now();
    let (result, _) = tokio::join!(
        client
            .converse()
            .model_id("m")
            .messages(user_message("Q"))
            .send(),
        serve_two,
    );

    assert!(result.is_ok(), "expected Ok but got: {:?}", result.err());
    started.elapsed()
}

// ── conformance §5.6 — 401 does NOT retry ────────────────────────────────────

#[tokio::test]
async fn no_retry_on_401() {
    let (listener, base_url) = bind().await;
    let attempt_count = Arc::new(AtomicUsize::new(0));
    let attempt_count2 = attempt_count.clone();

    let body = r#"{"error":{"message":"unauth","code":"invalid_api_key"}}"#;

    let client = Client::new(
        &Config::builder()
            .api_key("audacity_api_test")
            .base_url(&base_url)
            .max_retries(2)
            .build()
            .unwrap(),
    )
    .unwrap();

    // Serve only ONE response — if the SDK retries it would hang.
    let serve_task = async move {
        serve_once(&listener, http_response(401, "Unauthorized", "", body)).await;
        attempt_count2.fetch_add(1, Ordering::SeqCst);
    };

    let (result, _) = tokio::time::timeout(std::time::Duration::from_secs(5), async {
        tokio::join!(
            client
                .converse()
                .model_id("m")
                .messages(user_message("Q"))
                .send(),
            serve_task,
        )
    })
    .await
    .unwrap();

    assert!(matches!(result, Err(Error::AccessDenied(_))));
    assert_eq!(attempt_count.load(Ordering::SeqCst), 1);
}

// ── conformance §5.6 — BUDGET_EXCEEDED does NOT retry ────────────────────────

#[tokio::test]
async fn no_retry_on_budget_exceeded() {
    let (listener, base_url) = bind().await;

    let body = r#"{"success":false,"error":{"code":"BUDGET_EXCEEDED","message":"over budget"}}"#;

    let client = Client::new(
        &Config::builder()
            .api_key("audacity_api_test")
            .base_url(&base_url)
            .max_retries(2)
            .build()
            .unwrap(),
    )
    .unwrap();

    let serve_task = serve_once(&listener, http_response(429, "Too Many Requests", "", body));

    let (result, _) = tokio::time::timeout(std::time::Duration::from_secs(5), async {
        tokio::join!(
            client
                .converse()
                .model_id("m")
                .messages(user_message("Q"))
                .send(),
            serve_task,
        )
    })
    .await
    .unwrap();

    assert!(matches!(result, Err(Error::ServiceQuotaExceeded(_))));
}

// ── conformance §5.9 — missing API key ───────────────────────────────────────

#[test]
fn missing_api_key_fails_fast() {
    // Unset env var just in case. Safe despite the parallel test binary:
    // every other test passes api_key explicitly, so nothing else reads env.
    std::env::remove_var("AUDACITY_API_KEY");
    let result = Config::builder().build();
    assert!(matches!(result, Err(Error::MissingApiKey)));
}

// ── conformance §5.10 — User-Agent header ────────────────────────────────────

#[tokio::test]
async fn user_agent_header_sent() {
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let port = listener.local_addr().unwrap().port();
    let base_url = format!("http://127.0.0.1:{port}");

    let body = r#"{"choices":[{"message":{"content":"ok"},"finish_reason":"stop"}],"usage":{}}"#;
    let response = http_response(200, "OK", "", body);

    let capture_ua = async {
        let (mut sock, _) = listener.accept().await.unwrap();
        let req_str = read_request(&mut sock).await;
        sock.write_all(response.as_bytes()).await.unwrap();
        req_str
    };

    let (result, req_str) = tokio::join!(
        async {
            client_for(&base_url)
                .converse()
                .model_id("m")
                .messages(user_message("Q"))
                .send()
                .await
        },
        capture_ua,
    );

    result.unwrap();
    assert!(
        req_str.contains("audacity-sdk-rust/0.2.0"),
        "User-Agent not found in: {req_str}"
    );
}

// ── conformance §5.2 — streaming happy path ──────────────────────────────────

#[tokio::test]
async fn converse_stream_happy_path() {
    let (listener, base_url) = bind().await;

    let sse_body = concat!(
        "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"},\"index\":0}]}\n\n",
        "data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"},\"index\":0}]}\n\n",
        "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\",\"index\":0}]}\n\n",
        "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3}}\n\n",
        "data: [DONE]\n\n",
    );
    let response = http_sse_response(sse_body);

    let serve_task = serve_once(&listener, response);
    let stream_task = async {
        let mut out = client_for(&base_url)
            .converse_stream()
            .model_id("gpt-5.4-mini")
            .messages(user_message("Hi"))
            .send()
            .await?;

        let mut events: Vec<String> = Vec::new();
        while let Some(event) = out.stream.recv().await? {
            match &event {
                ConverseStreamOutput::MessageStart(_) => events.push("MessageStart".into()),
                ConverseStreamOutput::ContentBlockDelta(_) => {
                    events.push("ContentBlockDelta".into())
                }
                ConverseStreamOutput::ContentBlockStop(_) => events.push("ContentBlockStop".into()),
                ConverseStreamOutput::MessageStop(_) => events.push("MessageStop".into()),
                ConverseStreamOutput::Metadata(_) => events.push("Metadata".into()),
                _ => {}
            }
        }
        Ok::<_, Error>(events)
    };

    let (events, _) = tokio::join!(stream_task, serve_task);
    let events = events.unwrap();

    // Verify ordering: MessageStart → ContentBlockDelta → ContentBlockStop → MessageStop → Metadata
    assert_eq!(
        events,
        vec![
            "MessageStart",
            "ContentBlockDelta",
            "ContentBlockStop",
            "MessageStop",
            "Metadata",
        ]
    );
}

// ── conformance §5.3 — tool round-trip (streaming) ───────────────────────────

#[tokio::test]
async fn stream_tool_call_emits_content_block_start() {
    let (listener, base_url) = bind().await;

    let sse_body = concat!(
        "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\n",
        "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"type\":\"function\",\"function\":{\"name\":\"get_weather\",\"arguments\":\"\"}}]}}]}\n\n",
        "data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"{\\\"city\\\":\\\"NYC\\\"}\"}}]}}]}\n\n",
        "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"tool_calls\"}]}\n\n",
        "data: {\"choices\":[],\"usage\":{}}\n\n",
        "data: [DONE]\n\n",
    );
    let response = http_sse_response(sse_body);

    let stream_task = async {
        let mut out = client_for(&base_url)
            .converse_stream()
            .model_id("m")
            .messages(user_message("weather?"))
            .send()
            .await?;

        let mut has_start = false;
        let mut has_delta = false;
        let mut stop_reason = None;

        while let Some(event) = out.stream.recv().await? {
            match event {
                ConverseStreamOutput::ContentBlockStart(_) => has_start = true,
                ConverseStreamOutput::ContentBlockDelta(_) => has_delta = true,
                ConverseStreamOutput::MessageStop(e) => stop_reason = Some(e.stop_reason),
                _ => {}
            }
        }
        Ok::<_, Error>((has_start, has_delta, stop_reason))
    };

    let (result, _) = tokio::join!(stream_task, serve_once(&listener, response));
    let (has_start, has_delta, stop_reason) = result.unwrap();

    assert!(has_start, "expected ContentBlockStart for tool call");
    assert!(has_delta, "expected ContentBlockDelta for tool arguments");
    assert_eq!(stop_reason, Some(audacity_sdk::StopReason::ToolUse));
}

// ── chunked-transfer helpers for staged SSE writes ───────────────────────────

const SSE_HEADERS: &str =
    "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nTransfer-Encoding: chunked\r\n\r\n";

fn sse_chunk(data: &str) -> Vec<u8> {
    format!("{:x}\r\n{data}\r\n", data.len()).into_bytes()
}

const SSE_TERMINAL: &[u8] = b"0\r\n\r\n";

// ── spec §3 — connection drop before [DONE] → ModelStreamError ───────────────

#[tokio::test]
async fn stream_connection_drop_before_done_is_model_stream_error() {
    let (listener, base_url) = bind().await;

    let serve_task = async {
        let (mut sock, _) = listener.accept().await.unwrap();
        read_request(&mut sock).await;
        sock.write_all(SSE_HEADERS.as_bytes()).await.unwrap();
        sock.write_all(&sse_chunk(
            "data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\n",
        ))
        .await
        .unwrap();
        sock.flush().await.unwrap();
        // Drop the socket without the chunked terminator or [DONE].
        drop(sock);
    };

    let stream_task = async {
        let mut out = client_for(&base_url)
            .converse_stream()
            .model_id("m")
            .messages(user_message("Q"))
            .send()
            .await?;

        let mut saw_delta = false;
        loop {
            match out.stream.recv().await {
                Ok(Some(ConverseStreamOutput::ContentBlockDelta(_))) => saw_delta = true,
                Ok(Some(_)) => {}
                Ok(None) => return Ok::<_, Error>((saw_delta, None)),
                Err(e) => return Ok((saw_delta, Some(e))),
            }
        }
    };

    let (result, _) = tokio::join!(stream_task, serve_task);
    let (saw_delta, err) = result.unwrap();
    assert!(saw_delta, "expected the pre-drop delta to arrive");
    let err = err.expect("expected a stream error after the connection drop");
    assert!(
        matches!(err, Error::ModelStreamError(_)),
        "expected ModelStreamError, got {err:?}"
    );
    assert!(!err.is_retryable());
}

// ── spec §1/§4 — in-stream error chunk maps per the §4 code table ────────────

#[tokio::test]
async fn stream_inline_error_chunk_maps_per_code_table() {
    let (listener, base_url) = bind().await;

    let sse_body = concat!(
        "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\n",
        "data: {\"success\":false,\"error\":{\"code\":\"RATE_LIMIT_EXCEEDED\",\"message\":\"slow down\",\"request_id\":\"req-inline-1\"}}\n\n",
    );
    let response = http_sse_response(sse_body);

    let stream_task = async {
        let mut out = client_for(&base_url)
            .converse_stream()
            .model_id("m")
            .messages(user_message("Q"))
            .send()
            .await?;

        loop {
            match out.stream.recv().await {
                Ok(Some(_)) => {}
                Ok(None) => return Ok::<_, Error>(None),
                Err(e) => return Ok(Some(e)),
            }
        }
    };

    let (result, _) = tokio::join!(stream_task, serve_once(&listener, response));
    let err = result.unwrap().expect("expected an in-stream error");
    match err {
        Error::Throttling(d) => {
            assert_eq!(d.error_code.as_deref(), Some("RATE_LIMIT_EXCEEDED"));
            assert_eq!(d.request_id.as_deref(), Some("req-inline-1"));
            assert_eq!(d.status_code, 0);
        }
        other => panic!("expected Throttling (per §4 code table), got {other:?}"),
    }
}

// ── spec §4 — Retry-After lower-bounds the backoff sleep ─────────────────────

#[tokio::test]
async fn retry_after_lower_bounds_backoff() {
    let elapsed = converse_with_429_then_ok(1).await;
    assert!(
        elapsed >= std::time::Duration::from_secs(1),
        "backoff must sleep at least Retry-After (1s), only slept {elapsed:?}"
    );
}

// ── spec §1 — the timeout must NOT bound the SSE body read ───────────────────

#[tokio::test]
async fn stream_survives_past_configured_timeout() {
    let (listener, base_url) = bind().await;

    let client = Client::new(
        &Config::builder()
            .api_key("audacity_api_test")
            .base_url(&base_url)
            .timeout(std::time::Duration::from_secs(1))
            .max_retries(0)
            .build()
            .unwrap(),
    )
    .unwrap();

    let serve_task = async {
        let (mut sock, _) = listener.accept().await.unwrap();
        read_request(&mut sock).await;
        sock.write_all(SSE_HEADERS.as_bytes()).await.unwrap();
        sock.write_all(&sse_chunk(
            "data: {\"choices\":[{\"delta\":{\"content\":\"first\"}}]}\n\n",
        ))
        .await
        .unwrap();
        sock.flush().await.unwrap();
        // Stall well past the 1s configured timeout mid-body.
        tokio::time::sleep(std::time::Duration::from_secs(2)).await;
        sock.write_all(&sse_chunk(concat!(
            "data: {\"choices\":[{\"delta\":{\"content\":\"second\"},\"finish_reason\":null}]}\n\n",
            "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n",
            "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":2,\"total_tokens\":3}}\n\n",
            "data: [DONE]\n\n",
        )))
        .await
        .unwrap();
        sock.write_all(SSE_TERMINAL).await.unwrap();
    };

    let stream_task = async {
        let mut out = client
            .converse_stream()
            .model_id("m")
            .messages(user_message("Q"))
            .send()
            .await?;

        let mut deltas = 0;
        let mut saw_stop = false;
        let mut saw_metadata = false;
        while let Some(event) = out.stream.recv().await? {
            match event {
                ConverseStreamOutput::ContentBlockDelta(_) => deltas += 1,
                ConverseStreamOutput::MessageStop(_) => saw_stop = true,
                ConverseStreamOutput::Metadata(_) => saw_metadata = true,
                _ => {}
            }
        }
        Ok::<_, Error>((deltas, saw_stop, saw_metadata))
    };

    let (result, _) = tokio::join!(stream_task, serve_task);
    let (deltas, saw_stop, saw_metadata) =
        result.expect("stream must survive past the configured timeout");
    assert_eq!(deltas, 2);
    assert!(saw_stop);
    assert!(saw_metadata);
}