openlatch-client 0.1.18

//! Raw-socket HTTP/1.1 mock upstream for the boundary bench + streaming tests.
//!
//! `mockito` (the dev-dependency the rest of the suite uses) responds
//! instantly, which cannot discriminate a streaming forward from a buffering
//! one — against an instant mock both show ~simultaneous first/last byte
//! timestamps. The zero-buffer proof (C-9b) needs an upstream that emits chunks
//! with **real** inter-chunk delays, so this module stands up a tiny
//! `tokio::net::TcpListener` that writes chunked responses with `sleep` between
//! them. It is compiled only under `feature = "boundary"` and is used by the
//! shipped `bench` subcommands and the integration tests.

use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

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

/// A one-shot mock that captures the request and returns a fixed 200.
pub struct CaptureUpstream {
    /// Port the mock is listening on (loopback).
    pub port: u16,
    /// The exact request body bytes the mock received, once a request arrives.
    pub received_body: Arc<Mutex<Option<Vec<u8>>>>,
    /// The raw request header block (request line + headers) the mock received.
    pub received_headers: Arc<Mutex<Option<String>>>,
    /// The request-line + path the mock received (e.g. `POST /v1/messages HTTP/1.1`).
    pub received_request_line: Arc<Mutex<Option<String>>>,
}

impl CaptureUpstream {
    /// Case-insensitive lookup of a captured request header value.
    pub fn header(&self, name: &str) -> Option<String> {
        let headers = self.received_headers.lock().unwrap().clone()?;
        headers.lines().find_map(|l| {
            l.split_once(':')
                .filter(|(k, _)| k.eq_ignore_ascii_case(name))
                .map(|(_, v)| v.trim().to_string())
        })
    }
}

/// Read one HTTP/1.1 request off the socket: the header block, then the body —
/// decoding **either** a `Content-Length` body **or** a `Transfer-Encoding:
/// chunked` body (the boundary's opaque path forwards with chunked framing).
/// Returns `(header_block, body)`.
async fn read_request(stream: &mut TcpStream) -> std::io::Result<(String, Vec<u8>)> {
    let mut reader = BufReader::new(stream);

    // Header lines until the blank line.
    let mut header_block = String::new();
    loop {
        let mut line = String::new();
        let n = reader.read_line(&mut line).await?;
        if n == 0 {
            break;
        }
        let blank = line == "\r\n" || line == "\n";
        header_block.push_str(&line);
        if blank {
            break;
        }
    }

    let body = if let Some(cl) = parse_content_length(&header_block) {
        let mut buf = vec![0u8; cl];
        // A truncated body (client hangup) leaves a short buffer; tolerate it.
        let _ = reader.read_exact(&mut buf).await;
        buf
    } else if is_chunked(&header_block) {
        read_chunked(&mut reader).await?
    } else {
        Vec::new()
    };

    Ok((header_block, body))
}

/// Decode an HTTP/1.1 chunked body: `<hex-size>CRLF <data> CRLF …` until a
/// zero-length chunk.
async fn read_chunked(reader: &mut BufReader<&mut TcpStream>) -> std::io::Result<Vec<u8>> {
    let mut body = Vec::new();
    loop {
        let mut size_line = String::new();
        if reader.read_line(&mut size_line).await? == 0 {
            break;
        }
        let size = usize::from_str_radix(size_line.trim(), 16).unwrap_or(0);
        if size == 0 {
            // Consume the trailing CRLF after the terminating chunk.
            let mut trailer = String::new();
            let _ = reader.read_line(&mut trailer).await;
            break;
        }
        let mut chunk = vec![0u8; size];
        reader.read_exact(&mut chunk).await?;
        body.extend_from_slice(&chunk);
        // Consume the CRLF that follows the chunk data.
        let mut crlf = [0u8; 2];
        let _ = reader.read_exact(&mut crlf).await;
    }
    Ok(body)
}

fn parse_content_length(headers: &str) -> Option<usize> {
    headers
        .lines()
        .find_map(|l| {
            l.split_once(':')
                .filter(|(k, _)| k.eq_ignore_ascii_case("content-length"))
        })
        .and_then(|(_, v)| v.trim().parse().ok())
}

fn is_chunked(headers: &str) -> bool {
    headers.lines().any(|l| {
        l.split_once(':').is_some_and(|(k, v)| {
            k.eq_ignore_ascii_case("transfer-encoding")
                && v.to_ascii_lowercase().contains("chunked")
        })
    })
}

/// Spawn a mock upstream that captures the first request's body and replies
/// `200 OK` with the two-byte body `ok`. Returns immediately once bound.
pub async fn spawn_capture_200() -> CaptureUpstream {
    let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
    let port = listener.local_addr().unwrap().port();
    let received_body = Arc::new(Mutex::new(None));
    let received_headers = Arc::new(Mutex::new(None));
    let received_request_line = Arc::new(Mutex::new(None));
    let body_sink = received_body.clone();
    let hdr_sink = received_headers.clone();
    let line_sink = received_request_line.clone();

    tokio::spawn(async move {
        if let Ok((mut stream, _)) = listener.accept().await {
            if let Ok((headers, body)) = read_request(&mut stream).await {
                *line_sink.lock().unwrap() = headers.lines().next().map(str::to_string);
                *hdr_sink.lock().unwrap() = Some(headers);
                *body_sink.lock().unwrap() = Some(body);
            }
            let _ = stream
                .write_all(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok")
                .await;
            let _ = stream.flush().await;
        }
    });

    CaptureUpstream {
        port,
        received_body,
        received_headers,
        received_request_line,
    }
}

/// Handle onto a trickling SSE mock. `final_written_at` is set to the instant
/// the LAST chunk + terminator was flushed — the streaming test asserts the
/// client observed its first byte strictly before this.
pub struct TrickleUpstream {
    pub port: u16,
    pub final_written_at: Arc<Mutex<Option<Instant>>>,
}

/// Spawn a mock upstream that responds with `chunk_count` chunked-encoded SSE
/// events, sleeping `gap` between each. Records when the final chunk was
/// written so a test can prove the client saw byte one before the stream ended.
pub async fn spawn_trickle_sse(chunk_count: usize, gap: Duration) -> TrickleUpstream {
    let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
    let port = listener.local_addr().unwrap().port();
    let final_written_at = Arc::new(Mutex::new(None));
    let stamp = final_written_at.clone();

    tokio::spawn(async move {
        if let Ok((mut stream, _)) = listener.accept().await {
            let _ = read_request(&mut stream).await;
            let head = "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\
                        Transfer-Encoding: chunked\r\n\r\n";
            let _ = stream.write_all(head.as_bytes()).await;
            let _ = stream.flush().await;

            for i in 0..chunk_count {
                let payload = format!("data: chunk-{i}\n\n");
                let framed = format!("{:x}\r\n{}\r\n", payload.len(), payload);
                let _ = stream.write_all(framed.as_bytes()).await;
                let _ = stream.flush().await;
                if i + 1 < chunk_count {
                    tokio::time::sleep(gap).await;
                }
            }
            // Terminating zero-length chunk.
            let _ = stream.write_all(b"0\r\n\r\n").await;
            let _ = stream.flush().await;
            *stamp.lock().unwrap() = Some(Instant::now());
        }
    });

    TrickleUpstream {
        port,
        final_written_at,
    }
}

/// Spawn a mock upstream that captures the request body and replies with a
/// streaming SSE response whose terminal `message_delta` carries usage — the
/// happy path the capture tee reads (plan 02). The usage is parameterised so a
/// test can assert the exact token facts flow through (e.g. the C-3 shape).
///
/// Returns a [`CaptureUpstream`] so the test can also assert the forwarded body
/// (the two-sided privacy probe: the sentinel IS present here, in the bytes that
/// reached the provider).
#[allow(clippy::too_many_arguments)]
pub async fn spawn_capture_usage_sse(
    input_tokens: u64,
    cache_read: u64,
    cache_write: u64,
    eph_5m: u64,
    eph_1h: u64,
    output_tokens: u64,
) -> CaptureUpstream {
    let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
    let port = listener.local_addr().unwrap().port();
    let received_body = Arc::new(Mutex::new(None));
    let received_headers = Arc::new(Mutex::new(None));
    let received_request_line = Arc::new(Mutex::new(None));
    let body_sink = received_body.clone();
    let hdr_sink = received_headers.clone();
    let line_sink = received_request_line.clone();

    tokio::spawn(async move {
        if let Ok((mut stream, _)) = listener.accept().await {
            if let Ok((headers, body)) = read_request(&mut stream).await {
                *line_sink.lock().unwrap() = headers.lines().next().map(str::to_string);
                *hdr_sink.lock().unwrap() = Some(headers);
                *body_sink.lock().unwrap() = Some(body);
            }
            // message_start carries the input + cache fields (output_tokens = 1);
            // message_delta carries the final cumulative output — exactly how
            // Anthropic splits usage across a streamed turn.
            let sse = format!(
                "event: message_start\n\
                 data: {{\"type\":\"message_start\",\"message\":{{\"usage\":{{\"input_tokens\":{input_tokens},\"cache_read_input_tokens\":{cache_read},\"cache_creation_input_tokens\":{cache_write},\"cache_creation\":{{\"ephemeral_5m_input_tokens\":{eph_5m},\"ephemeral_1h_input_tokens\":{eph_1h}}},\"output_tokens\":1}}}}}}\n\n\
                 event: content_block_delta\n\
                 data: {{\"type\":\"content_block_delta\",\"delta\":{{\"text\":\"hello\"}}}}\n\n\
                 event: message_delta\n\
                 data: {{\"type\":\"message_delta\",\"usage\":{{\"output_tokens\":{output_tokens}}}}}\n\n"
            );
            let head = format!(
                "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\
                 Content-Length: {}\r\nConnection: close\r\n\r\n",
                sse.len()
            );
            let _ = stream.write_all(head.as_bytes()).await;
            let _ = stream.write_all(sse.as_bytes()).await;
            let _ = stream.flush().await;
        }
    });

    CaptureUpstream {
        port,
        received_body,
        received_headers,
        received_request_line,
    }
}

/// Spawn a mock upstream that captures the request body and replies with a 200
/// SSE response carrying ONLY the `message_start` event — final input/cache but
/// the PRELIMINARY `output_tokens: 1` — and then CLOSES the connection WITHOUT
/// ever sending the terminal `message_delta`. Exercises FIX 1: the stream ends
/// before terminal usage, so the emitted event must degrade to a local estimate
/// (`tokenizer_estimated` / `stream_interrupted`), never `provider_reported`
/// with the preliminary output.
///
/// The body is framed by connection-close (no Content-Length, no chunked), so
/// the forwarder observes a clean stream end after the `message_start` bytes.
pub async fn spawn_message_start_then_hangup() -> CaptureUpstream {
    let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
    let port = listener.local_addr().unwrap().port();
    let received_body = Arc::new(Mutex::new(None));
    let received_headers = Arc::new(Mutex::new(None));
    let received_request_line = Arc::new(Mutex::new(None));
    let body_sink = received_body.clone();
    let hdr_sink = received_headers.clone();
    let line_sink = received_request_line.clone();

    tokio::spawn(async move {
        if let Ok((mut stream, _)) = listener.accept().await {
            if let Ok((headers, body)) = read_request(&mut stream).await {
                *line_sink.lock().unwrap() = headers.lines().next().map(str::to_string);
                *hdr_sink.lock().unwrap() = Some(headers);
                *body_sink.lock().unwrap() = Some(body);
            }
            // message_start ONLY — final input/cache, preliminary output=1.
            let start = "event: message_start\n\
                 data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":40,\"cache_read_input_tokens\":0,\"cache_creation_input_tokens\":0,\"output_tokens\":1}}}\n\n";
            // No Content-Length, no chunked → body delimited by connection close.
            let head =
                "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nConnection: close\r\n\r\n";
            let _ = stream.write_all(head.as_bytes()).await;
            let _ = stream.write_all(start.as_bytes()).await;
            let _ = stream.flush().await;
            // Drop the socket WITHOUT a message_delta → the stream ends here.
            drop(stream);
        }
    });

    CaptureUpstream {
        port,
        received_body,
        received_headers,
        received_request_line,
    }
}

/// A guaranteed-closed loopback port: bind then drop, returning the freed port
/// number. Used to force a connect-refused for the synthetic-502 posture test.
pub async fn closed_port() -> u16 {
    let l = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
    let p = l.local_addr().unwrap().port();
    drop(l);
    p
}

/// Spawn a mock upstream that ACCEPTS the connection and reads the request, then
/// holds the socket open WITHOUT ever sending response headers. Exercises the
/// header-wait timeout (FIX 1): `connect_timeout` is satisfied (the TCP/TLS
/// handshake completes), but the forwarder's `.send()` would otherwise block
/// forever waiting for a status line that never comes. Returns the port; the
/// spawned task owns the connection and sleeps well past any test-scale timeout.
pub async fn spawn_hang_after_accept() -> u16 {
    let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
    let port = listener.local_addr().unwrap().port();

    tokio::spawn(async move {
        if let Ok((mut stream, _)) = listener.accept().await {
            // Drain the request so the client finishes writing, then go silent:
            // no status line, no headers. Hold the socket so it is not a
            // connection reset (which would be a different failure path).
            let _ = read_request(&mut stream).await;
            tokio::time::sleep(Duration::from_secs(30)).await;
            drop(stream);
        }
    });

    port
}