openlatch-client 0.5.4

//! Raw-socket HTTP/1.1 mock upstream for the model relay 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 = "model-relay"` 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 model relay'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
/// model relay 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 SSE turn into the shape that broke the round-5 scanner: the `lead`
/// events, then the terminal `frame` in **three** pieces with a **newline-free
/// middle** one, then an optional `trail` (a `[DONE]` sentinel, where the format
/// has 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.
///
/// A chat-completions usage chunk is small enough to arrive whole in practice,
/// and it is split here anyway: the carry-over is shared code, and a format that
/// only ever exercises the easy path is a format the shared code can be broken
/// under without anything going red.
fn sse_pieces(lead: &str, frame: &str, trail: &str) -> Vec<String> {
    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 mut pieces = vec![
        lead.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"
    );
    if !trail.is_empty() {
        pieces.push(trail.to_string());
    }
    pieces
}

/// The lead event every Responses split is preceded by.
const RESPONSES_DELTA_EVENT: &str = "event: response.output_text.delta\n\
                                     data: {\"type\":\"response.output_text.delta\",\"delta\":\"hi\"}\n\n";

/// A **real captured** pair of chat-completions content chunks — the same
/// recording [`super::capture::CHAT_COMPLETIONS_CHUNK_FIXTURE`] came from, so
/// the stream this mock serves is the stream that was observed rather than a
/// belief about it. `ollama 0.34.0`, `qwen2.5-coder:14b`, 2026-09-14.
const CHAT_COMPLETIONS_DELTA_EVENTS: &str = concat!(
    r#"data: {"id":"chatcmpl-398","object":"chat.completion.chunk","created":1789410009,"model":"qwen2.5-coder:14b","system_fingerprint":"fp_ollama","choices":[{"index":0,"delta":{"role":"assistant","content":"hi"},"finish_reason":null}]}"#,
    "\n\n",
    r#"data: {"id":"chatcmpl-398","object":"chat.completion.chunk","created":1789410009,"model":"qwen2.5-coder:14b","system_fingerprint":"fp_ollama","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}"#,
    "\n\n",
);

/// The sentinel that ends a chat-completions stream, after the usage chunk.
const CHAT_COMPLETIONS_DONE: &str = "data: [DONE]\n\n";

/// One chat-completions **terminal usage chunk**, framed as an SSE event,
/// carrying the given counts.
///
/// The chunk is [`super::capture::CHAT_COMPLETIONS_CHUNK_FIXTURE`] — a recorded
/// capture, 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. The EMPTY `choices` array the
/// recording carries is preserved, which is what makes a proxy-level test
/// evidence about the real terminal chunk.
fn chat_completions_usage_chunk(
    prompt_tokens: u64,
    cached_tokens: u64,
    completion_tokens: u64,
) -> String {
    let total = prompt_tokens + completion_tokens;
    // The same member order the captured chunk uses.
    let usage = format!(
        "\"usage\":{{\"prompt_tokens\":{prompt_tokens},\
         \"prompt_tokens_details\":{{\"cached_tokens\":{cached_tokens}}},\
         \"completion_tokens\":{completion_tokens},\
         \"total_tokens\":{total}}}"
    );
    let body = super::capture::CHAT_COMPLETIONS_CHUNK_FIXTURE
        .replace(super::capture::CHAT_COMPLETIONS_CHUNK_FIXTURE_USAGE, &usage);
    assert!(
        body.contains(&usage),
        "the fixture's usage object must be substitutable — the two consts have drifted"
    );
    format!("data: {body}\n\n")
}

/// 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 {
    spawn_capture_sse_pieces(sse_pieces(
        RESPONSES_DELTA_EVENT,
        &responses_completed_frame(
            input_tokens,
            cached_tokens,
            cache_write_tokens,
            output_tokens,
        ),
        "",
    ))
    .await
}

/// Spawn a mock upstream that captures the request body and replies with an
/// **OpenAI chat-completions** SSE turn whose terminal chunk carries the given
/// counts — the `stream_options: {include_usage: true}` branch of DD-07.
///
/// Split the same way, for the same reason: a terminal chunk arriving in pieces
/// with a newline-free middle one is what D-15's carry-over exists for, and
/// sending it whole would let a scanner that drops its held tail through.
pub async fn spawn_capture_chat_completions_sse(
    prompt_tokens: u64,
    cached_tokens: u64,
    completion_tokens: u64,
) -> CaptureUpstream {
    spawn_capture_sse_pieces(sse_pieces(
        CHAT_COMPLETIONS_DELTA_EVENTS,
        &chat_completions_usage_chunk(prompt_tokens, cached_tokens, completion_tokens),
        CHAT_COMPLETIONS_DONE,
    ))
    .await
}

/// One **generate-content** chunk, framed as an SSE event, carrying the given
/// counts.
///
/// The chunk is [`super::capture::GOOGLE_GENERATE_CONTENT_FIXTURE`] with only
/// its `usageMetadata` object replaced, so the mock and the unit fixture are the
/// same payload and a change to one cannot silently diverge from the other.
/// `finishReason` is whatever the fixture carries — `"STOP"`, i.e. terminal.
fn google_usage_chunk(
    prompt_tokens: u64,
    cached_tokens: u64,
    candidates_tokens: u64,
    thoughts_tokens: u64,
) -> String {
    // Google's own identity: prompt + candidates + thoughts (this mock sends no
    // tool-use tokens), with `cachedContentTokenCount` a SUBSET of the prompt.
    let total = prompt_tokens + candidates_tokens + thoughts_tokens;
    // The same member order the fixture uses.
    let usage = format!(
        "\"usageMetadata\":{{\"promptTokenCount\":{prompt_tokens},\
         \"cachedContentTokenCount\":{cached_tokens},\
         \"candidatesTokenCount\":{candidates_tokens},\
         \"thoughtsTokenCount\":{thoughts_tokens},\
         \"toolUsePromptTokenCount\":0,\
         \"totalTokenCount\":{total}}}"
    );
    let body = super::capture::GOOGLE_GENERATE_CONTENT_FIXTURE.replace(
        super::capture::GOOGLE_GENERATE_CONTENT_FIXTURE_USAGE,
        &usage,
    );
    assert!(
        body.contains(&usage),
        "the fixture's usageMetadata object must be substitutable — the two consts have drifted"
    );
    format!("data: {body}\n\n")
}

/// A non-terminal generate-content chunk: content, and a `usageMetadata` whose
/// counts are SMALLER than the terminal chunk's, with **no `finishReason`**.
///
/// This is the lead every Google split is preceded by, and it is not
/// decoration. Gemini puts `usageMetadata` on every chunk and grows it as the
/// stream runs, so a decoder keying "terminal" on the presence of usage marks
/// THIS chunk terminal and an accumulator that SUMS bills the turn twice. Both
/// mistakes are visible only when the lead carries usage too.
///
/// Its INPUT-side counts (`promptTokenCount`, `cachedContentTokenCount`,
/// `toolUsePromptTokenCount`) match what [`spawn_capture_google_sse`] puts on
/// the terminal chunk, because Gemini holds them constant across a turn — only
/// the output side grows. A lead disagreeing with the terminal chunk would make
/// the element-wise max pick the LEAD's input, and the test asserting the
/// subtraction would be asserting the wrong chunk.
const GOOGLE_DELTA_EVENT: &str = concat!(
    r#"data: {"candidates":[{"content":{"parts":[{"text":"The parser "}],"role":"model"},"index":0}],"#,
    r#""usageMetadata":{"promptTokenCount":1024,"cachedContentTokenCount":256,"candidatesTokenCount":9,"thoughtsTokenCount":4,"toolUsePromptTokenCount":0,"totalTokenCount":1037},"#,
    r#""modelVersion":"gemini-2.5-pro","responseId":"y8nHaPeaDvqfmecPmLXzsAo"}"#,
    "\n\n",
);

/// Spawn a mock upstream that captures the request and replies with a
/// **generate-content** `alt=sse` turn: a non-terminal chunk that already
/// carries (smaller) cumulative usage, then the terminal chunk carrying
/// `finishReason` and the final counts.
///
/// Split the same way the other two are, and for one extra reason: the cut is
/// asserted to land INSIDE the `usageMetadata` object, so a scanner that
/// dropped its held tail would read no usage at all rather than a slightly
/// wrong number.
pub async fn spawn_capture_google_sse(
    prompt_tokens: u64,
    cached_tokens: u64,
    candidates_tokens: u64,
    thoughts_tokens: u64,
) -> CaptureUpstream {
    let pieces = sse_pieces(
        GOOGLE_DELTA_EVENT,
        &google_usage_chunk(
            prompt_tokens,
            cached_tokens,
            candidates_tokens,
            thoughts_tokens,
        ),
        "",
    );
    // MEASURED, not assumed: no piece of the TERMINAL chunk may hold the whole
    // usage object, or this mock exercises the easy path and proves nothing
    // about D-15. `pieces[0]` is the lead event, whose own usageMetadata is
    // deliberately whole — it is there to carry cumulative counts, not to be cut.
    assert!(
        !pieces[1..]
            .iter()
            .any(|p| p.contains("\"usageMetadata\"") && p.contains("\"totalTokenCount\"")),
        "the terminal usageMetadata must straddle a chunk boundary"
    );
    spawn_capture_sse_pieces(pieces).await
}

/// Spawn a mock upstream replying with a **complete, well-formed**
/// chat-completions turn that carries **no usage object at all** — the other
/// branch of DD-07.
///
/// This is what a client that does not send `stream_options.include_usage`
/// receives, and we may not add the flag to its request: `AGENTS.md` makes
/// provider-traffic mutation default-off, and everything fails toward the
/// original bytes. The stream below is the real recording of that case
/// (`ollama 0.34.0`, 2026-09-14) — content chunks, a `finish_reason`, the
/// sentinel, and the literal string "usage" nowhere in it.
pub async fn spawn_capture_chat_completions_sse_without_usage() -> CaptureUpstream {
    let stream = format!("{CHAT_COMPLETIONS_DELTA_EVENTS}{CHAT_COMPLETIONS_DONE}");
    assert!(
        !stream.contains("usage"),
        "the no-usage branch must carry no usage object — otherwise it tests the other branch"
    );
    spawn_capture_sse_pieces(stream.split_inclusive("\n\n").map(str::to_string).collect()).await
}

/// Spawn a mock upstream replying with an **Ollama native** `/api/chat` stream:
/// newline-delimited JSON with no `data:` framing, the lines a live
/// `qwen2.5-coder:7b` sends.
///
/// The last content line and the terminal `done: true` line arrive in ONE
/// chunk, on purpose: a terminal line sent alone decodes through the scanner's
/// whole-body fallback whether its line gate is right or not, so only the
/// coalesced shape proves the gate.
pub async fn spawn_capture_ollama_native_stream(
    prompt_eval_count: u64,
    prompt_eval_cached_count: u64,
    eval_count: u64,
) -> CaptureUpstream {
    let content = |text: &str| {
        format!(
            "{{\"model\":\"qwen2.5-coder:7b\",\"message\":{{\"role\":\"assistant\",\"content\":\"{text}\"}},\"done\":false}}\n"
        )
    };
    let done = format!(
        "{{\"model\":\"qwen2.5-coder:7b\",\"message\":{{\"role\":\"assistant\",\"content\":\"\"}},\
         \"done\":true,\"done_reason\":\"stop\",\"prompt_eval_count\":{prompt_eval_count},\
         \"prompt_eval_cached_count\":{prompt_eval_cached_count},\"eval_count\":{eval_count}}}\n"
    );
    spawn_capture_sse_pieces(vec![content("p"), format!("{}{done}", content("ong"))]).await
}

/// The socket half every piece-wise SSE mock shares: capture the request, then
/// write each piece as its own HTTP chunk with a flush and a real gap, so the
/// forwarder observes them as separate chunks rather than one coalesced body.
async fn spawn_capture_sse_pieces(pieces: Vec<String>) -> 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);
            }
            // 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**
/// model relay, 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 }
}

/// Capture two requests, rejecting the first with 400 and accepting the retry.
pub async fn spawn_reject_then_capture_usage_sse() -> MultiCaptureUpstream {
    spawn_error_then_capture_usage_sse("400 Bad Request").await
}

/// Capture two requests, returning one chosen error and then a usage-bearing 200.
pub async fn spawn_error_then_capture_usage_sse(status: &'static str) -> 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 {
        let mut count = 0usize;
        while let Ok((mut stream, _)) = listener.accept().await {
            if let Ok((_headers, body)) = read_request(&mut stream).await {
                sink.lock().unwrap().push(body);
            }
            count += 1;
            let response = if count == 1 {
                format!("HTTP/1.1 {status}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")
            } else {
                let sse = "event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":1}}\n\n";
                format!("HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{sse}", sse.len())
            };
            let _ = stream.write_all(response.as_bytes()).await;
            let _ = stream.flush().await;
            if count >= 2 {
                break;
            }
        }
    });
    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, drawn below every platform's ephemeral range.
/// Used to force a connect-refused for the synthetic-502 posture tests.
pub async fn closed_port() -> u16 {
    use std::net::{SocketAddr, TcpStream as StdTcpStream};
    use std::sync::atomic::{AtomicU16, Ordering};

    static NEXT: AtomicU16 = AtomicU16::new(20000);

    loop {
        let port = NEXT.fetch_add(1, Ordering::Relaxed);
        assert!(port < 21000, "exhausted 20000..21000 dead ports");
        let addr = SocketAddr::from(([127, 0, 0, 1], port));
        if StdTcpStream::connect_timeout(&addr, Duration::from_millis(50)).is_err() {
            return port;
        }
    }
}

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