openlatch-client 0.3.3

//! 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,
    }
}

/// Spawn a mock upstream that answers **every** connection `200 OK`, forever.
///
/// [`spawn_capture_200`] accepts exactly one connection, which is right for a
/// test that asserts on the one request it made and wrong for any test that
/// probes more than once — the second probe gets a connection reset, the
/// boundary answers its own synthetic 502, and the test reads a *provider*
/// failure where the code under test was fine. The wiring-loop tests probe once
/// per agent, so they need this one.
///
/// Returns the bound loopback port; nothing is captured, because a test that
/// needs the request body wants the one-shot capture instead.
pub async fn spawn_always_200() -> u16 {
    let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
    let port = listener.local_addr().unwrap().port();
    tokio::spawn(async move {
        while let Ok((mut stream, _)) = listener.accept().await {
            tokio::spawn(async move {
                if read_request(&mut stream).await.is_ok() {
                    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;
                }
            });
        }
    });
    port
}

/// 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,
    }
}

/// One **real** `response.completed` frame with the usage counts substituted.
///
/// The frame is [`super::capture::RESPONSES_COMPLETED_FIXTURE`] — a recorded
/// capture of the OpenAI Responses API, not a hand-written shape — with only
/// its `usage` object replaced, so the mock and the unit fixture are the same
/// payload and a change to one cannot silently diverge from the other.
fn responses_completed_frame(
    input_tokens: u64,
    cached_tokens: u64,
    cache_write_tokens: u64,
    output_tokens: u64,
) -> String {
    let total = input_tokens + output_tokens;
    // The same member order the captured frame uses.
    let usage = format!(
        "\"usage\":{{\"input_tokens\":{input_tokens},\
         \"input_tokens_details\":{{\"cache_write_tokens\":{cache_write_tokens},\"cached_tokens\":{cached_tokens}}},\
         \"output_tokens\":{output_tokens},\
         \"output_tokens_details\":{{\"reasoning_tokens\":0}},\
         \"total_tokens\":{total}}}"
    );
    let body = super::capture::RESPONSES_COMPLETED_FIXTURE
        .replace(super::capture::RESPONSES_COMPLETED_FIXTURE_USAGE, &usage);
    assert!(
        body.contains(&usage),
        "the fixture's usage object must be substitutable — the two consts have drifted"
    );
    format!("event: response.completed\ndata: {body}\n\n")
}

/// Cut one Responses turn into the shape that broke the round-5 scanner: a
/// leading delta event, then the terminal frame in **three** pieces with a
/// **newline-free middle** one.
///
/// This is not decoration. A `response.completed` frame embeds the whole
/// `Response` object, so on a live turn it arrives in four or more chunks and
/// its middle chunks carry no newline at all — a scanner that drops its held
/// tail on a newline-free chunk never reassembles the frame, and the turn
/// degrades on every request. Sending the frame whole here would let that bug
/// through the proxy-level tests untouched.
fn responses_sse_pieces(frame: &str) -> Vec<String> {
    let delta = "event: response.output_text.delta\n\
                 data: {\"type\":\"response.output_text.delta\",\"delta\":\"hi\"}\n\n";
    let json_at = frame
        .find("data: ")
        .expect("the frame carries one data: line")
        + "data: ".len();
    let a = json_at + 40;
    let b = json_at + (frame.len() - json_at) / 2;
    let pieces = vec![
        delta.to_string(),
        frame[..a].to_string(),
        frame[a..b].to_string(),
        frame[b..].to_string(),
    ];
    assert!(
        !pieces[2].contains('\n'),
        "the middle piece must be newline-free — it is the case D-15 exists for"
    );
    pieces
}

/// Spawn a mock upstream that captures the request body and replies with an
/// **OpenAI Responses** SSE turn whose terminal `response.completed` frame
/// carries the given usage counts.
///
/// Every other SSE mock in this module is Anthropic-shaped by construction
/// (`spawn_capture_usage_sse` formats `message_start` / `message_delta`), so
/// none of them can drive the Responses decoder.
///
/// The frame is sent in **three chunks with a newline-free middle one**, after
/// a leading delta event, so the proxy-level tests exercise D-15's carry-over
/// through the real tee rather than only through the unit scanner.
pub async fn spawn_capture_responses_sse(
    input_tokens: u64,
    cached_tokens: u64,
    cache_write_tokens: 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();

    let pieces = responses_sse_pieces(&responses_completed_frame(
        input_tokens,
        cached_tokens,
        cache_write_tokens,
        output_tokens,
    ));

    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);
            }
            // Chunked, one HTTP chunk per piece with a flush and a real gap
            // between them, so the forwarder observes them as separate chunks
            // rather than one coalesced body.
            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 piece in &pieces {
                let framed = format!("{:x}\r\n{}\r\n", piece.len(), piece);
                let _ = stream.write_all(framed.as_bytes()).await;
                let _ = stream.flush().await;
                tokio::time::sleep(Duration::from_millis(5)).await;
            }
            let _ = stream.write_all(b"0\r\n\r\n").await;
            let _ = stream.flush().await;
        }
    });

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

/// Spawn a mock upstream that serves **many** requests in sequence, recording
/// every request body in order.
///
/// The single-shot [`spawn_capture_usage_sse`] accepts exactly one connection,
/// which is the right shape for a one-request assertion and the wrong one for
/// anything that has to observe a *session*. L-0 is measured over turns — its
/// horizon is repetition already billed for — so proving it acts (and proving it
/// does not act with the flag off) needs several requests through **one**
/// boundary, against one upstream, with each forwarded body kept.
///
/// `bodies` collects them in arrival order.
pub struct MultiCaptureUpstream {
    /// Port the mock is listening on (loopback).
    pub port: u16,
    /// Every request body received, in arrival order.
    pub bodies: Arc<Mutex<Vec<Vec<u8>>>>,
}

impl MultiCaptureUpstream {
    /// Wait until at least `n` requests have arrived, or give up after ~3s.
    pub async fn wait_for(&self, n: usize) -> Vec<Vec<u8>> {
        for _ in 0..300 {
            {
                let got = self.bodies.lock().unwrap();
                if got.len() >= n {
                    return got.clone();
                }
            }
            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
        }
        self.bodies.lock().unwrap().clone()
    }
}

/// Spawn a [`MultiCaptureUpstream`] replying to every request with the same
/// usage-bearing SSE turn.
pub async fn spawn_multi_capture_usage_sse(
    input_tokens: u64,
    output_tokens: u64,
) -> MultiCaptureUpstream {
    let listener = TcpListener::bind(("127.0.0.1", 0)).await.unwrap();
    let port = listener.local_addr().unwrap().port();
    let bodies = Arc::new(Mutex::new(Vec::new()));
    let sink = bodies.clone();

    tokio::spawn(async move {
        while let Ok((mut stream, _)) = listener.accept().await {
            if let Ok((_headers, body)) = read_request(&mut stream).await {
                sink.lock().unwrap().push(body);
            }
            let sse = format!(
                "event: message_start\n\
                 data: {{\"type\":\"message_start\",\"message\":{{\"usage\":{{\"input_tokens\":{input_tokens},\"cache_read_input_tokens\":0,\"cache_creation_input_tokens\":0,\"output_tokens\":1}}}}}}\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;
        }
    });

    MultiCaptureUpstream { port, bodies }
}

/// 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
}