AkurAI Build
Menu

AkurAI-Build

public

Latest change 266a1aa21e346770b48e96f008ab6c2a5a603961 - Complete PR interfaces, isolated CI and crash-safe merge recovery by Ólafur Búi Ólafsson

#![allow(dead_code)]

use std::{
    collections::BTreeMap,
    fmt::Display,
    io::{BufRead, BufReader, Read, Write},
    net::TcpStream,
    path::Path,
    process::{Child, Command, Stdio},
    sync::{Arc, Mutex},
    thread,
    time::{Duration, Instant},
};

fn result_or_panic<T, E: Display>(result: Result<T, E>, context: &str) -> T {
    result.unwrap_or_else(|error| panic!("{context}: {error}"))
}

fn option_or_panic<T>(value: Option<T>, context: &str) -> T {
    value.unwrap_or_else(|| panic!("{context}"))
}
pub struct ProcessServer {
    child: Option<Child>,
    address: String,
    logs: Arc<Mutex<String>>,
}

impl ProcessServer {
    /// Start the server bound to an OS-assigned loopback port (`--address
    /// 127.0.0.1:0`) and learn the real port from the child's own startup log.
    /// This is race-free: the kernel assigns a unique free port at bind time,
    /// so parallel tests can never collide on a port the way a
    /// reserve-then-release scheme can.
    pub fn start_dynamic(mut command: Command) -> Self {
        command
            .stdin(Stdio::null())
            .stdout(Stdio::piped())
            .stderr(Stdio::piped());
        let mut child = result_or_panic(command.spawn(), "spawn AkurAI Build server");
        let logs = Arc::new(Mutex::new(String::new()));
        if let Some(out) = child.stdout.take() {
            spawn_log_reader(out, Arc::clone(&logs));
        }
        if let Some(err) = child.stderr.take() {
            spawn_log_reader(err, Arc::clone(&logs));
        }
        let deadline = Instant::now() + Duration::from_secs(15);
        loop {
            if let Some(address) = listen_address(&logs) {
                return Self {
                    child: Some(child),
                    address,
                    logs,
                };
            }
            if let Some(status) = result_or_panic(child.try_wait(), "poll AkurAI Build server") {
                panic!(
                    "AkurAI Build exited before listening ({status}):\n{}",
                    snapshot(&logs)
                );
            }
            if Instant::now() >= deadline {
                let _ = child.kill();
                let _ = child.wait();
                panic!(
                    "AkurAI Build did not report a listening address:\n{}",
                    snapshot(&logs)
                );
            }
            thread::sleep(Duration::from_millis(100));
        }
    }

    pub fn address(&self) -> &str {
        &self.address
    }

    pub fn request(&self, method: &str, target: &str) -> HttpResponse {
        self.request_with_body(method, target, &[], &[])
    }

    pub fn request_with_headers(
        &self,
        method: &str,
        target: &str,
        headers: &[(&str, &str)],
    ) -> HttpResponse {
        self.request_with_body(method, target, headers, &[])
    }

    pub fn request_with_body(
        &self,
        method: &str,
        target: &str,
        headers: &[(&str, &str)],
        body: &[u8],
    ) -> HttpResponse {
        let mut stream = result_or_panic(
            TcpStream::connect(&self.address),
            "connect to bunfork server",
        );
        result_or_panic(
            stream.set_read_timeout(Some(Duration::from_secs(5))),
            "set HTTP read timeout",
        );
        result_or_panic(
            stream.set_write_timeout(Some(Duration::from_secs(5))),
            "set HTTP write timeout",
        );
        let mut request = format!("{method} {target} HTTP/1.1\r\nHost: {}\r\n", self.address);
        for (name, value) in headers {
            request.push_str(name);
            request.push_str(": ");
            request.push_str(value);
            request.push_str("\r\n");
        }
        request.push_str(&format!(
            "Content-Length: {}\r\nConnection: close\r\n\r\n",
            body.len()
        ));
        result_or_panic(
            stream.write_all(request.as_bytes()),
            "write HTTP request head",
        );
        result_or_panic(stream.write_all(body), "write HTTP request body");
        let mut bytes = Vec::new();
        result_or_panic(stream.read_to_end(&mut bytes), "read HTTP response");
        HttpResponse::parse(&bytes)
    }

    pub fn shutdown_with_sigterm(&mut self) {
        let logs = Arc::clone(&self.logs);
        let child = option_or_panic(self.child.as_mut(), "server process missing");
        let signal = result_or_panic(
            Command::new("kill")
                .args(["-TERM", &child.id().to_string()])
                .status(),
            "send SIGTERM",
        );
        assert!(signal.success(), "kill -TERM failed: {signal}");

        let deadline = Instant::now() + Duration::from_secs(5);
        loop {
            if let Some(status) = result_or_panic(child.try_wait(), "wait for SIGTERM shutdown") {
                assert!(
                    status.success(),
                    "bunfork did not exit cleanly after SIGTERM ({status}):\n{}",
                    snapshot(&logs)
                );
                self.child = None;
                return;
            }
            if Instant::now() >= deadline {
                let _ = child.kill();
                let _ = child.wait();
                panic!("bunfork did not stop after SIGTERM:\n{}", snapshot(&logs));
            }
            thread::sleep(Duration::from_millis(20));
        }
    }
}

impl Drop for ProcessServer {
    fn drop(&mut self) {
        if let Some(child) = &mut self.child {
            let _ = child.kill();
            let _ = child.wait();
        }
    }
}

#[derive(Debug)]
pub struct HttpResponse {
    pub status: u16,
    pub headers: BTreeMap<String, String>,
    pub body: Vec<u8>,
}

impl HttpResponse {
    fn parse(bytes: &[u8]) -> Self {
        let separator = option_or_panic(
            bytes.windows(4).position(|window| window == b"\r\n\r\n"),
            "HTTP response omitted header terminator",
        );
        let head = result_or_panic(
            std::str::from_utf8(&bytes[..separator]),
            "HTTP headers are not UTF-8",
        );
        let mut lines = head.split("\r\n");
        let status = option_or_panic(
            lines
                .next()
                .and_then(|line| line.split_whitespace().nth(1))
                .and_then(|value| value.parse().ok()),
            "HTTP response omitted status",
        );
        let mut headers = BTreeMap::<String, String>::new();
        for line in lines {
            let (name, value) = option_or_panic(line.split_once(':'), "malformed HTTP header");
            headers
                .entry(name.trim().to_ascii_lowercase())
                .and_modify(|existing| {
                    existing.push_str(", ");
                    existing.push_str(value.trim());
                })
                .or_insert_with(|| value.trim().to_owned());
        }
        let raw_body = &bytes[separator + 4..];
        let body = if raw_body.is_empty() {
            Vec::new()
        } else if headers
            .get("transfer-encoding")
            .is_some_and(|value| value.eq_ignore_ascii_case("chunked"))
        {
            decode_chunked(raw_body)
        } else {
            raw_body.to_vec()
        };
        Self {
            status,
            headers,
            body,
        }
    }

    pub fn assert_status(&self, expected: u16) {
        assert_eq!(self.status, expected, "unexpected HTTP response: {self:?}");
    }

    pub fn assert_header_contains(&self, name: &str, expected: &str) {
        let value = self
            .headers
            .get(name)
            .unwrap_or_else(|| panic!("missing {name} header: {self:?}"));
        assert!(
            value
                .to_ascii_lowercase()
                .contains(&expected.to_ascii_lowercase()),
            "{name}={value:?} does not contain {expected:?}"
        );
    }

    pub fn assert_body_contains(&self, expected: &str) {
        let body = String::from_utf8_lossy(&self.body);
        assert!(
            body.contains(expected),
            "body does not contain {expected:?}: {body:?}"
        );
    }

    pub fn assert_body_excludes(&self, unexpected: &str) {
        let body = String::from_utf8_lossy(&self.body);
        assert!(
            !body.contains(unexpected),
            "body unexpectedly contains {unexpected:?}: {body:?}"
        );
    }

    pub fn assert_security_headers(&self) {
        self.assert_header_contains("x-content-type-options", "nosniff");
        self.assert_header_contains("x-frame-options", "deny");
        self.assert_header_contains("referrer-policy", "strict-origin-when-cross-origin");
        self.assert_header_contains("content-security-policy", "default-src 'self'");
        self.assert_header_contains("permissions-policy", "camera=()");
    }
}

fn decode_chunked(bytes: &[u8]) -> Vec<u8> {
    let mut decoded = Vec::new();
    let mut cursor = 0;
    loop {
        let line_end = option_or_panic(
            bytes[cursor..]
                .windows(2)
                .position(|window| window == b"\r\n")
                .map(|offset| cursor + offset),
            "invalid chunk size line",
        );
        let chunk_line = result_or_panic(
            std::str::from_utf8(&bytes[cursor..line_end]),
            "chunk size is not ASCII",
        );
        let size_text = option_or_panic(chunk_line.split(';').next(), "empty chunk size");
        let size = result_or_panic(usize::from_str_radix(size_text, 16), "invalid chunk size");
        cursor = line_end + 2;
        if size == 0 {
            break;
        }
        let end = option_or_panic(cursor.checked_add(size), "chunk size overflow");
        assert!(end + 2 <= bytes.len(), "truncated chunked response");
        decoded.extend_from_slice(&bytes[cursor..end]);
        assert_eq!(&bytes[end..end + 2], b"\r\n", "invalid chunk terminator");
        cursor = end + 2;
    }
    decoded
}

/// Snapshot the accumulated child log output without disturbing the readers.
fn snapshot(logs: &Arc<Mutex<String>>) -> String {
    logs.lock().map(|guard| guard.clone()).unwrap_or_default()
}

fn listen_address(logs: &Arc<Mutex<String>>) -> Option<String> {
    logs.lock().ok()?.lines().find_map(parse_listen_addr)
}

/// Extract the loopback socket used by this harness from a tracing line.
/// The literal address remains intact even when ambient settings add ANSI decoration.
fn parse_listen_addr(line: &str) -> Option<String> {
    let start = line.find("127.0.0.1:")?;
    let address: String = line[start..]
        .chars()
        .take_while(|character| character.is_ascii_digit() || matches!(character, '.' | ':'))
        .collect();
    address.parse::<std::net::SocketAddr>().ok()?;
    Some(address)
}

/// Drain a child pipe into the shared log buffer.
fn spawn_log_reader<R: Read + Send + 'static>(reader: R, logs: Arc<Mutex<String>>) {
    thread::spawn(move || {
        let buffered = BufReader::new(reader);
        for line in buffered.lines() {
            let Ok(line) = line else { break };
            if let Ok(mut guard) = logs.lock() {
                guard.push_str(&line);
                guard.push('\n');
            }
        }
    });
}

pub struct ScratchDir(tempfile::TempDir);

impl ScratchDir {
    pub fn new(label: &str) -> Self {
        let directory = result_or_panic(
            tempfile::Builder::new()
                .prefix(&format!("bunfork-{label}-"))
                .tempdir(),
            "create test scratch directory",
        );
        Self(directory)
    }

    pub fn path(&self) -> &Path {
        self.0.path()
    }
}