newt-core 0.8.0

//! Primary streaming dispatch, distinct from the final display reissue.
use super::*;
use crate::agentic::anthropic_loop_tests::FixtureMcpPermission;
use crate::test_guard::{hang_guarded, HANG_GUARD};
use serde_json::{json, Value};
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};

const TOOL: &str = "stream_fixture__read_result";

#[derive(Default)]
struct FixtureTools(Vec<String>);

#[async_trait::async_trait]
impl McpTools for FixtureTools {
    fn handles(&self, name: &str) -> bool {
        name == TOOL
    }
    fn tool_defs(&self) -> Vec<Value> {
        vec![json!({"type":"function","function":{"name":TOOL,
            "description":"Read a deterministic fixture.","parameters":{"type":"object",
            "properties":{"stage":{"type":"string"}},"required":["stage"]}}})]
    }
    async fn call(&mut self, call: &LeasedMcpCall<'_>) -> String {
        let stage = call.args()["stage"].as_str().unwrap().to_owned();
        self.0.push(stage.clone());
        format!("fixture result {stage}")
    }
}

fn wire(frames: &[Value], done: bool) -> ResponseTemplate {
    let mut body = frames
        .iter()
        .map(|frame| format!("data: {frame}\n\n"))
        .collect::<String>();
    if done {
        body.push_str("data: [DONE]\n\n");
    }
    ResponseTemplate::new(200).set_body_raw(body.into_bytes(), "text/event-stream")
}

fn streamed_batch(done: bool, valid_arguments: bool) -> ResponseTemplate {
    streamed_batch_finishing(done, valid_arguments, "tool_calls")
}

/// The two-call batch, ending with `finish`. With `valid_arguments` false, the
/// second call's arguments are cut mid-string.
fn streamed_batch_finishing(done: bool, valid_arguments: bool, finish: &str) -> ResponseTemplate {
    wire(
        &[
            json!({"model":"served-model","choices":[{"delta":{"reasoning_content":"fixture reasoning",
            "tool_calls":[
                {"index":0,"id":"call-a","type":"function","function":{"name":TOOL,"arguments":"{\"stage\":"}},
                {"index":1,"id":"call-b","type":"function","function":{"name":TOOL,"arguments":"{\"stage\":"}}
            ]}}]}),
            json!({"choices":[{"delta":{"tool_calls":[
            {"index":0,"function":{"arguments":"\"a\"}"}},
            {"index":1,"function":{"arguments":if valid_arguments { "\"b\"}" } else { "\"b" }}}
        ]},"finish_reason":finish}]}),
            json!({"choices":[],"usage":{"prompt_tokens":2000,"completion_tokens":4}}),
        ],
        done,
    )
}

fn streamed_answer() -> ResponseTemplate {
    wire(
        &[
            json!({"choices":[{"delta":{"content":"Fixture response."},"finish_reason":"stop"}]}),
            json!({"choices":[],"usage":{"prompt_tokens":2100,"completion_tokens":2}}),
        ],
        true,
    )
}

fn complete_single_tool_frame() -> Value {
    json!({"model":"served-model","choices":[{"index":0,"delta":{"tool_calls":[{
        "index":0,"id":"call-a","type":"function","function":{
            "name":TOOL,"arguments":"{\"stage\":\"a\"}"
        }
    }]},"finish_reason":"tool_calls"}]})
}

fn complete_single_tool_prefix() -> String {
    format!("data: {}\n\n", complete_single_tool_frame())
}

/// #2372: an accepted final answer is generated ONCE and returned exactly as
/// the gates accepted it. The mock would answer any second request with
/// DIFFERENT text, so a display reissue shows up twice: as an extra request,
/// and as `REISSUED` in the returned reply.
#[tokio::test]
async fn an_accepted_final_answer_is_generated_once_and_returned_as_accepted() {
    let server = MockServer::start().await;
    let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0));
    let counter = calls.clone();
    Mock::given(method("POST"))
        .and(path("/v1/chat/completions"))
        .respond_with(move |_request: &Request| {
            let text = if counter.fetch_add(1, Ordering::SeqCst) == 0 {
                "ACCEPTED answer."
            } else {
                "REISSUED answer."
            };
            wire(
                &[
                    json!({"choices":[{"delta":{"content":text},"finish_reason":"stop"}]}),
                    json!({"choices":[],"usage":{"prompt_tokens":1000,"completion_tokens":3}}),
                ],
                true,
            )
        })
        .mount(&server)
        .await;
    let uri = server.uri();
    let messages = msgs();
    let caveats = Caveats::top();
    let mut context = ctx(&uri, &messages, &caveats);
    context.action_nudges = false;
    let (text, streamed, usage, _) = chat_complete(context, &mut NoMcp)
        .await
        .expect("a plain answer completes the turn");
    assert_eq!(
        crate::agentic::model_explanation(&text),
        "ACCEPTED answer.",
        "the gated answer, never a second generation"
    );
    assert!(!streamed, "nothing was streamed to a display");
    assert_eq!(
        server.received_requests().await.unwrap().len(),
        1,
        "one generation request"
    );
    assert_eq!(usage.unwrap().output_tokens, 3, "one generation's tokens");
}

#[tokio::test]
async fn primary_stream_assembles_a_tool_batch_once_and_preserves_call_ids() {
    let server = MockServer::start().await;
    let seen = Arc::new(Mutex::new(Vec::<Value>::new()));
    let requests = seen.clone();
    Mock::given(method("POST"))
        .and(path("/v1/chat/completions"))
        .respond_with(move |request: &Request| {
            let body = body_json(request);
            let has_result = body["messages"]
                .as_array()
                .unwrap()
                .iter()
                .any(|message| message["role"] == "tool");
            requests.lock().unwrap().push(body);
            if has_result {
                streamed_answer()
            } else {
                streamed_batch(true, true)
            }
        })
        .mount(&server)
        .await;
    let uri = server.uri();
    let messages = msgs();
    let caveats = Caveats::top();
    let mut context = ctx(&uri, &messages, &caveats);
    context.action_nudges = false;
    let mut permission = FixtureMcpPermission::new(TOOL, &caveats);
    context.permission_gate = Some(&mut permission);
    let mut tools = FixtureTools::default();
    let (text, streamed, usage, hallucinations) = chat_complete(context, &mut tools)
        .await
        .expect("a complete streamed tool batch must reach execution");
    assert_eq!(tools.0, ["a", "b"]);
    assert_eq!(
        crate::agentic::model_explanation(&text),
        "Fixture response."
    );
    assert!(!streamed, "the host renders the accepted answer (#2372)");
    assert_eq!(hallucinations, 0);
    assert_eq!(usage.unwrap().output_tokens, 6, "tool batch 4 + answer 2");
    let seen = seen.lock().unwrap();
    assert_eq!(seen.len(), 2, "tool batch, final answer");
    for request in seen.iter() {
        assert_eq!(request["stream"], true);
        assert_eq!(request["stream_options"]["include_usage"], true);
    }
    let results = seen[1]["messages"]
        .as_array()
        .unwrap()
        .iter()
        .filter(|message| message["role"] == "tool")
        .collect::<Vec<_>>();
    assert_eq!(results.len(), 2);
    assert_eq!(results[0]["tool_call_id"], "call-a");
    assert_eq!(results[0]["content"], "fixture result a");
    assert_eq!(results[1]["tool_call_id"], "call-b");
    assert_eq!(results[1]["content"], "fixture result b");
}

#[tokio::test]
async fn primary_stream_rejects_a_cut_batch_before_any_tool_runs() {
    let server = MockServer::start().await;
    Mock::given(method("POST"))
        .and(path("/v1/chat/completions"))
        .respond_with(streamed_batch(false, true))
        .mount(&server)
        .await;
    let uri = server.uri();
    let messages = msgs();
    let caveats = Caveats::top();
    let mut context = ctx(&uri, &messages, &caveats);
    context.action_nudges = false;
    let mut permission = FixtureMcpPermission::new(TOOL, &caveats);
    context.permission_gate = Some(&mut permission);
    let mut tools = FixtureTools::default();
    assert!(chat_complete(context, &mut tools).await.is_err());
    assert!(tools.0.is_empty(), "no partial batch may authorize a tool");
    let requests = server.received_requests().await.unwrap();
    let generation = requests
        .iter()
        .filter(|request| request.url.path() == "/v1/chat/completions")
        .collect::<Vec<_>>();
    assert_eq!(generation.len(), 1);
    assert_eq!(body_json(generation[0])["stream"], true);
}

/// The same two-call batch as a complete JSON reply, which the strict decoder
/// passes through as-is: the non-streamed Chat twin of [`streamed_batch_finishing`].
fn json_batch(finish: &str) -> ResponseTemplate {
    ResponseTemplate::new(200).set_body_json(json!({"model":"served-model",
        "choices":[{"index":0,"message":{"role":"assistant","content":"","tool_calls":[
            {"id":"call-a","type":"function","function":{"name":TOOL,"arguments":"{\"stage\":\"a\"}"}},
            {"id":"call-b","type":"function","function":{"name":TOOL,"arguments":"{\"stage\":\"b"}}
        ]},"finish_reason":finish}],
        "usage":{"prompt_tokens":2000,"completion_tokens":4}}))
}

/// #2385: a complete stream whose tool call carries invalid arguments may be
/// retried when the provider did not report an output limit (#2824). The loop echoes a keyed rejection per call and asks again, as the
/// non-streamed Chat wire does; before, the strict decoder failed the turn. No
/// call of the rejected batch runs, the retry's batch runs once, and both wires
/// send the model the same rejection.
#[tokio::test]
async fn invalid_streamed_arguments_are_rejected_per_call_and_retried() {
    let mut rejections = Vec::new();
    for (streamed, finish) in [(true, "tool_calls"), (false, "tool_calls")] {
        let label = format!("streamed={streamed} finish={finish}");
        let server = MockServer::start().await;
        let seen = Arc::new(Mutex::new(Vec::<Value>::new()));
        let requests = seen.clone();
        let finish_owned = finish.to_owned();
        Mock::given(method("POST"))
            .and(path("/v1/chat/completions"))
            .respond_with(move |request: &Request| {
                let body = body_json(request);
                let tool_messages = body["messages"]
                    .as_array()
                    .unwrap()
                    .iter()
                    .filter(|message| message["role"] == "tool")
                    .count();
                requests.lock().unwrap().push(body);
                match tool_messages {
                    0 if streamed => streamed_batch_finishing(true, false, &finish_owned),
                    0 => json_batch(&finish_owned),
                    2 => streamed_batch(true, true),
                    _ => streamed_answer(),
                }
            })
            .mount(&server)
            .await;
        let uri = server.uri();
        let messages = msgs();
        let caveats = Caveats::top();
        let mut context = ctx(&uri, &messages, &caveats);
        context.action_nudges = false;
        let mut permission = FixtureMcpPermission::new(TOOL, &caveats);
        context.permission_gate = Some(&mut permission);
        let mut tools = FixtureTools::default();
        let (text, _, _, _) = chat_complete(context, &mut tools)
            .await
            .unwrap_or_else(|error| panic!("{label}: the turn continues: {error:#}"));
        assert_eq!(
            crate::agentic::model_explanation(&text),
            "Fixture response.",
            "{label}"
        );
        assert_eq!(tools.0, ["a", "b"], "{label}: only the retried batch runs");
        let seen = seen.lock().unwrap();
        assert_eq!(
            seen.len(),
            3,
            "{label}: rejected batch, retried batch, answer (#2388 removed the display reissue)"
        );
        let echoed: Vec<_> = seen[1]["messages"]
            .as_array()
            .unwrap()
            .iter()
            .filter(|message| message["role"] == "tool")
            .map(|message| {
                (
                    message["tool_call_id"].clone(),
                    message["content"].as_str().unwrap().to_owned(),
                )
            })
            .collect();
        assert_eq!(
            echoed.iter().map(|(id, _)| id.clone()).collect::<Vec<_>>(),
            [json!("call-a"), json!("call-b")],
            "{label}: one keyed rejection per call"
        );
        assert!(
            echoed.iter().all(
                |(_, content)| content.starts_with("tool-call batch rejected before execution")
            ),
            "{label}: {echoed:?}"
        );
        rejections.push(echoed);
    }
    assert!(
        rejections.windows(2).all(|pair| pair[0] == pair[1]),
        "every wire and finish sends the same rejection: {rejections:#?}"
    );
}

#[tokio::test]
async fn cancelled_stream_with_a_complete_tool_call_dispatches_nothing() {
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let uri = format!("http://{}", listener.local_addr().unwrap());
    let (delivered_tx, delivered_rx) = tokio::sync::oneshot::channel();
    let server = tokio::spawn(async move {
        let (mut socket, _) = listener.accept().await.unwrap();
        let (_, request) = read_request(&mut socket).await;
        socket
            .write_all(
                format!(
                    "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n\
                     Connection: close\r\n\r\n{}",
                    complete_single_tool_prefix()
                )
                .as_bytes(),
            )
            .await
            .unwrap();
        let padding = format!(":{}\n", "x".repeat(64 * 1024 - 2));
        for _ in 0..256 {
            socket.write_all(padding.as_bytes()).await.unwrap();
        }
        delivered_tx.send(request).unwrap();
        match socket.read(&mut [0_u8; 1]).await {
            Ok(0) => {}
            Err(error)
                if matches!(
                    error.kind(),
                    std::io::ErrorKind::ConnectionReset | std::io::ErrorKind::ConnectionAborted
                ) => {}
            other => panic!("cancelled generation socket remained active: {other:?}"),
        }
    });

    let messages = msgs();
    let caveats = Caveats::top();
    let cancel = AtomicBool::new(false);
    let mut context = ctx(&uri, &messages, &caveats);
    context.action_nudges = false;
    context.cancel = Some(&cancel);
    let mut permission = FixtureMcpPermission::new(TOOL, &caveats);
    context.permission_gate = Some(&mut permission);
    let mut tools = FixtureTools::default();
    let mut completion = Box::pin(chat_complete(context, &mut tools));
    let request = hang_guarded("the complete tool-call frame reaching the client", async {
        tokio::select! {
            request = delivered_rx => request.unwrap(),
            result = &mut completion => panic!("unfinished response resolved: {result:?}"),
        }
    })
    .await;
    assert_eq!(request["stream"], true);

    cancel.store(true, Ordering::Relaxed);
    let (reply, _, _, _) = hang_guarded("cancellation dropping the response body", &mut completion)
        .await
        .expect("cancellation is a clean turn outcome");
    assert!(reply.is_empty());
    drop(completion);
    server.await.unwrap();
    assert_eq!(
        tools.0.len(),
        0,
        "a complete call is inert until the whole streamed response is accepted"
    );
}

#[tokio::test]
async fn completed_stream_with_the_same_tool_call_dispatches_once() {
    let server = MockServer::start().await;
    let round = Arc::new(std::sync::atomic::AtomicUsize::new(0));
    Mock::given(method("POST"))
        .and(path("/v1/chat/completions"))
        .respond_with({
            let round = round.clone();
            move |_: &Request| match round.fetch_add(1, Ordering::Relaxed) {
                0 => wire(&[complete_single_tool_frame()], true),
                _ => streamed_answer(),
            }
        })
        .mount(&server)
        .await;
    let messages = msgs();
    let caveats = Caveats::top();
    let uri = server.uri();
    let mut context = ctx(&uri, &messages, &caveats);
    context.action_nudges = false;
    let mut permission = FixtureMcpPermission::new(TOOL, &caveats);
    context.permission_gate = Some(&mut permission);
    let mut tools = FixtureTools::default();

    chat_complete(context, &mut tools)
        .await
        .expect("the same call followed by [DONE] must be accepted");

    assert_eq!(tools.0, ["a"], "the completed twin must exercise dispatch");
}

#[tokio::test]
async fn provider_error_after_a_complete_tool_call_dispatches_nothing() {
    let server = MockServer::start().await;
    let body = format!(
        "{}data: {{\"error\":{{\"message\":\"fixture rejection\",\"code\":400}}}}\n\n",
        complete_single_tool_prefix()
    );
    Mock::given(method("POST"))
        .and(path("/v1/chat/completions"))
        .respond_with(ResponseTemplate::new(200).set_body_raw(body, "text/event-stream"))
        .mount(&server)
        .await;
    let messages = msgs();
    let caveats = Caveats::top();
    let uri = server.uri();
    let mut context = ctx(&uri, &messages, &caveats);
    context.action_nudges = false;
    let mut permission = FixtureMcpPermission::new(TOOL, &caveats);
    context.permission_gate = Some(&mut permission);
    let mut tools = FixtureTools::default();

    assert!(chat_complete(context, &mut tools).await.is_err());
    assert_eq!(
        tools.0.len(),
        0,
        "an observed provider rejection cannot authorize a preceding call"
    );
}

#[tokio::test]
async fn cap_summary_stream_decodes_the_answer_and_usage_without_tools() {
    let server = MockServer::start().await;
    Mock::given(method("POST"))
        .and(path("/v1/chat/completions"))
        .respond_with(streamed_answer())
        .mount(&server)
        .await;
    let uri = server.uri();
    let messages = msgs();
    let caveats = Caveats::top();
    let mut context = ctx(&uri, &messages, &caveats);
    context.action_nudges = false;
    context.max_tool_rounds = 0;
    let (text, streamed, usage, _) = chat_complete(context, &mut NoMcp).await.unwrap();
    assert_eq!(
        crate::agentic::model_explanation(&text),
        "Fixture response."
    );
    assert!(
        !streamed,
        "the cap summary is returned for its existing display path"
    );
    assert_eq!(usage.unwrap().output_tokens, 2);
    let requests = server.received_requests().await.unwrap();
    let generation = requests
        .iter()
        .filter(|request| request.url.path() == "/v1/chat/completions")
        .collect::<Vec<_>>();
    assert_eq!(generation.len(), 1);
    let body = body_json(generation[0]);
    assert_eq!(body["stream"], true);
    assert_eq!(body["stream_options"]["include_usage"], true);
    assert!(body.get("tools").is_none());
}

async fn progressing_core_stream_resets_its_idle_timeout(cap_summary: bool) {
    const FRAMES: [&[u8]; 4] = [
        b"data: {\"choices\":[{\"index\":0,\"delta\":{\"role\":\"assistant\"}}]}\n\n",
        b"data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"progressing \"}}]}\n\n",
        b"data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"answer\"}}]}\n\n",
        b"data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\ndata: [DONE]\n\n",
    ];
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let uri = format!("http://{}", listener.local_addr().unwrap());
    let (started_tx, started_rx) = tokio::sync::oneshot::channel();
    let (advance_tx, mut advance_rx) = tokio::sync::mpsc::unbounded_channel();
    let server = tokio::spawn(async move {
        let (mut socket, _) = listener.accept().await.unwrap();
        let (path, request) = read_request(&mut socket).await;
        assert_eq!(path, "/v1/chat/completions");
        socket
            .write_all(
                b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nConnection: close\r\n\r\n",
            )
            .await
            .unwrap();
        socket.write_all(FRAMES[0]).await.unwrap();
        started_tx.send(request.clone()).unwrap();
        for frame in &FRAMES[1..] {
            advance_rx.recv().await.expect("test advances the stream");
            socket.write_all(frame).await.unwrap();
        }
        socket.shutdown().await.unwrap();
        request
    });

    let messages = msgs();
    let caveats = Caveats::top();
    let mut context = ctx(&uri, &messages, &caveats);
    context.action_nudges = false;
    context.inference_timeout_secs = 1;
    if cap_summary {
        context.max_tool_rounds = 0;
    }
    let mut tools = NoMcp;
    let consumed = Arc::new(std::sync::atomic::AtomicUsize::new(0));
    let mut completion = Box::pin(
        crate::retry::RESPONSE_BYTES_READ
            .scope(consumed.clone(), chat_complete(context, &mut tools)),
    );
    let request = hang_guarded("the generation beginning", async {
        tokio::select! {
            request = started_rx => request.unwrap(),
            result = &mut completion => panic!("stream ended before its first frame: {result:?}"),
        }
    })
    .await;
    assert_eq!(request["stream"], true);

    tokio::time::pause();
    // An independent REAL-clock watchdog for the loop below. `hang_guarded`
    // (a `tokio::time::timeout`) measures against THIS runtime's clock, which
    // is paused and only advances via `advance()` or when the runtime goes
    // idle — and the loop below stays runnable on purpose (`yield_now`), to
    // stop paused time auto-advancing while the real socket catches up. A
    // stalled reader or a `consumed` mismatch would therefore spin forever
    // without ever tripping a paused-clock timeout. A plain OS thread sleeping
    // the REAL `HANG_GUARD` is immune to the pause and can still catch a
    // genuine hang; the loop polls the flag it sets, same shape as
    // `hang_guarded`, just driven by a real clock instead of a paused one.
    let watchdog_fired = Arc::new(AtomicBool::new(false));
    let watchdog_done = Arc::new(AtomicBool::new(false));
    {
        let fired = Arc::clone(&watchdog_fired);
        let done = Arc::clone(&watchdog_done);
        std::thread::spawn(move || {
            std::thread::sleep(HANG_GUARD);
            if !done.load(Ordering::SeqCst) {
                fired.store(true, Ordering::SeqCst);
            }
        });
    }
    let mut expected = 0;
    let mut completed = None;
    for (index, frame) in FRAMES.iter().enumerate() {
        if index > 0 {
            tokio::time::advance(Duration::from_millis(400)).await;
            advance_tx.send(()).unwrap();
        }
        expected += frame.len();
        // A server write does not prove the client reset its idle timer. Drive
        // the actual body reader until it consumes this frame before advancing
        // time again. Yield stays runnable, preventing paused-time auto-advance
        // while the real socket's readiness notification catches up.
        while consumed.load(Ordering::SeqCst) < expected || index == FRAMES.len() - 1 {
            assert!(
                !watchdog_fired.load(Ordering::SeqCst),
                "hung: the progress loop produced no result within the {}s \
                 real-clock hang guard",
                HANG_GUARD.as_secs()
            );
            tokio::select! {
                biased;
                result = &mut completion => {
                    assert_eq!(index, FRAMES.len() - 1,
                        "stream ended before its terminal frame: {result:?}");
                    completed = Some(result);
                    break;
                }
                _ = tokio::task::yield_now() => {}
            }
        }
        assert_eq!(consumed.load(Ordering::SeqCst), expected);
    }
    watchdog_done.store(true, Ordering::SeqCst);
    tokio::time::resume();
    let result = completed.expect("the progressing generation must finish");
    let (text, streamed, _, _) = result.expect("progress must reset the idle timeout");
    server.await.unwrap();
    // #2372: both the primary answer and the cap summary are returned as
    // generated; nothing is sent again for display.
    assert_eq!(
        crate::agentic::model_explanation(&text),
        "progressing answer"
    );
    assert!(!streamed);
}

#[tokio::test]
async fn primary_progressing_stream_resets_its_idle_timeout() {
    progressing_core_stream_resets_its_idle_timeout(false).await;
}

#[tokio::test]
async fn cap_summary_progressing_stream_resets_its_idle_timeout() {
    progressing_core_stream_resets_its_idle_timeout(true).await;
}

async fn read_request(socket: &mut tokio::net::TcpStream) -> (String, Value) {
    let mut bytes = Vec::new();
    let mut chunk = [0_u8; 4096];
    let header_end = loop {
        let read = socket.read(&mut chunk).await.unwrap();
        assert!(read > 0);
        bytes.extend_from_slice(&chunk[..read]);
        if let Some(end) = bytes.windows(4).position(|window| window == b"\r\n\r\n") {
            break end + 4;
        }
    };
    let headers = std::str::from_utf8(&bytes[..header_end]).unwrap();
    let path = headers
        .lines()
        .next()
        .unwrap()
        .split_whitespace()
        .nth(1)
        .unwrap()
        .to_owned();
    let length: usize = headers
        .lines()
        .find_map(|line| {
            let (name, value) = line.split_once(':')?;
            name.eq_ignore_ascii_case("content-length")
                .then(|| value.trim().parse().unwrap())
        })
        .unwrap();
    while bytes.len() - header_end < length {
        let read = socket.read(&mut chunk).await.unwrap();
        assert!(read > 0);
        bytes.extend_from_slice(&chunk[..read]);
    }
    (
        path,
        serde_json::from_slice(&bytes[header_end..header_end + length]).unwrap(),
    )
}

/// Grounds the decoder tests' cancellation assumption in an actual HTTP peer.
/// Counts may be observed first; the generation blocks behind a deterministic
/// channel so cancellation cannot race the count or its matching dispatch.
async fn cancellation_case(set_flag: bool, cap_summary: bool, measured: bool) {
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let uri = format!("http://{}", listener.local_addr().unwrap());
    let (started, reading) = tokio::sync::oneshot::channel();
    let server = tokio::spawn(async move {
        let mut measurement: Option<(Value, usize, usize)> = None;
        let (mut socket, body) = loop {
            let (mut socket, _) = listener.accept().await.unwrap();
            let (path, body) = read_request(&mut socket).await;
            if path == "/v1/chat/completions" {
                if measured {
                    let (counted, _, _) = measurement.as_ref().expect("count precedes generation");
                    assert_eq!(
                        &body, counted,
                        "generation must match the complete counted request"
                    );
                }
                break (socket, body);
            }
            if measured && path == "/v1/chat/completions/input_tokens" {
                let prior = estimate_request_tokens(
                    body["messages"].as_array().unwrap(),
                    body.get("tools"),
                    crate::tokens::TokenEstimation::default(),
                );
                let tokens = prior * 2;
                assert!(
                    tokens <= 100_000,
                    "fixture count must fit its configured input bound"
                );
                assert!(measurement.replace((body, prior, tokens)).is_none());
                let response = json!({"input_tokens":tokens}).to_string();
                socket.write_all(format!(
                    "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{response}", response.len()
                ).as_bytes()).await.unwrap();
                socket.shutdown().await.unwrap();
                continue;
            }
            socket
                .write_all(
                    b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
                )
                .await
                .unwrap();
            socket.shutdown().await.unwrap();
        };
        socket.write_all(b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nConnection: close\r\n\r\ndata: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\n").await.unwrap();
        started.send(body).unwrap();
        match socket.read(&mut [0_u8; 1]).await {
            Ok(0) => (),
            Err(error)
                if matches!(
                    error.kind(),
                    std::io::ErrorKind::ConnectionReset | std::io::ErrorKind::ConnectionAborted
                ) => {}
            other => panic!("the generation socket remained active: {other:?}"),
        }
        measurement.map(|(_, prior, tokens)| (prior, tokens))
    });
    let messages = msgs();
    let caveats = Caveats::top();
    let cancel = AtomicBool::new(false);
    let mut state = CompressState::new();
    let mut context = ctx(&uri, &messages, &caveats);
    context.action_nudges = false;
    context.cancel = Some(&cancel);
    context.compress_state = Some(&mut state);
    if measured {
        context.safe_context = Some(100_000);
    }
    if cap_summary {
        context.max_tool_rounds = 0;
    }
    let mut tools = NoMcp;
    let mut completion = Box::pin(chat_complete(context, &mut tools));
    let request = hang_guarded("the generation reaching the server", async {
        tokio::select! {
            request = reading => request.unwrap(),
            result = &mut completion => panic!("unfinished response resolved: {result:?}"),
        }
    })
    .await;
    let result = if set_flag {
        cancel.store(true, Ordering::Relaxed);
        Some(hang_guarded("the cancel flag interrupting the read", &mut completion).await)
    } else {
        None
    };
    drop(completion);
    let measurement = hang_guarded("the server observing socket closure", server)
        .await
        .unwrap();
    // Clean up before asserting on either streaming generation path.
    if let Some(result) = result {
        let (reply, _, _, _) = result.unwrap();
        assert!(reply.is_empty());
    }
    assert_eq!(request["stream"], true);
    assert_eq!(request["stream_options"]["include_usage"], true);
    if let Some((prior, tokens)) = measurement {
        assert_eq!(
            state.calibration.ratio(None),
            tokens as f32 / prior as f32,
            "cancellation must retain the completed count of the same request"
        );
    }
}

#[tokio::test]
async fn primary_stream_dropping_future_closes_the_generation_socket() {
    cancellation_case(false, false, false).await;
}

#[tokio::test]
async fn primary_stream_cancel_flag_closes_the_generation_socket() {
    cancellation_case(true, false, false).await;
}

#[tokio::test]
async fn cap_summary_stream_cancel_flag_closes_the_generation_socket() {
    cancellation_case(true, true, false).await;
}

#[tokio::test]
async fn primary_stream_cancellation_retains_its_completed_token_count() {
    cancellation_case(true, false, true).await;
}

#[tokio::test]
async fn cap_summary_stream_cancellation_retains_its_completed_token_count() {
    cancellation_case(true, true, true).await;
}

/// Grounds the response reader's status precedence in a real incomplete HTTP
/// frame: a known provider status is not replaced by a later socket failure.
#[tokio::test]
async fn provider_http_status_survives_a_truncated_response_body() {
    use crate::retry::Retryability;
    for (status, expected) in [
        ("401 Unauthorized", Retryability::Fatal),
        ("503 Service Unavailable", Retryability::Retry),
    ] {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let uri = format!("http://{}", listener.local_addr().unwrap());
        let server = tokio::spawn(async move {
            let (mut socket, _) = listener.accept().await.unwrap();
            read_request(&mut socket).await;
            socket.write_all(format!("HTTP/1.1 {status}\r\nContent-Length: 128\r\nConnection: close\r\n\r\nfixture provider failure").as_bytes()).await.unwrap();
            socket.shutdown().await.unwrap();
        });
        let result = hang_guarded("the truncated body ending the read", async {
            let response = reqwest::Client::new()
                .post(uri)
                .json(&json!({}))
                .send()
                .await
                .unwrap();
            crate::agentic::smart_harness::response(response, None, "inference endpoint").await
        })
        .await;
        server.await.unwrap();
        let error = result.unwrap_err();
        assert_eq!(crate::retry::classify(&error), expected, "{error}");
        assert_eq!(
            crate::agentic::observability::error_class(&error),
            Some(crate::agentic::observability::ErrorClass::Model)
        );
        assert!(
            error.to_string().contains("fixture provider failure"),
            "{error}"
        );
        assert!(
            error.to_string().contains("response body read failed"),
            "{error}"
        );
    }
}

#[tokio::test]
async fn successful_stream_mid_body_disconnect_is_retryable() {
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let uri = format!("http://{}", listener.local_addr().unwrap());
    let server = tokio::spawn(async move {
        let (mut socket, _) = listener.accept().await.unwrap();
        read_request(&mut socket).await;
        socket
            .write_all(
                b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: 4096\r\nConnection: close\r\n\r\ndata: {\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n\n",
            )
            .await
            .unwrap();
        socket.shutdown().await.unwrap();
    });
    let response = reqwest::Client::new()
        .post(uri)
        .json(&json!({}))
        .send()
        .await
        .unwrap();
    let error = crate::agentic::smart_harness::response_with_decoder(
        response,
        None,
        "inference endpoint",
        crate::agentic::smart_harness::decode_openai_response,
    )
    .await
    .unwrap_err();
    server.await.unwrap();
    assert_eq!(
        crate::retry::classify(&error),
        crate::retry::Retryability::Retry,
        "a cut successful stream is a retryable transport failure: {error:#}"
    );
    assert_eq!(
        crate::agentic::observability::error_class(&error),
        Some(crate::agentic::observability::ErrorClass::Transport),
        "a body cut is a transport failure even when the status was successful"
    );
}

#[tokio::test]
async fn successful_stream_idle_timeout_is_retryable_and_typed_timeout() {
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let uri = format!("http://{}", listener.local_addr().unwrap());
    let server = tokio::spawn(async move {
        let (mut socket, _) = listener.accept().await.unwrap();
        read_request(&mut socket).await;
        socket
            .write_all(
                b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nConnection: close\r\n\r\ndata: {\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n\n",
            )
            .await
            .unwrap();
        match socket.read(&mut [0_u8; 1]).await {
            Ok(0) => {}
            Err(error)
                if matches!(
                    error.kind(),
                    std::io::ErrorKind::ConnectionReset | std::io::ErrorKind::ConnectionAborted
                ) => {}
            other => panic!("idle timeout left the response socket active: {other:?}"),
        }
    });
    let response = reqwest::Client::builder()
        .read_timeout(Duration::from_millis(100))
        .build()
        .unwrap()
        .post(uri)
        .json(&json!({}))
        .send()
        .await
        .unwrap();
    let error = crate::agentic::smart_harness::response_with_decoder(
        response,
        None,
        "inference endpoint",
        crate::agentic::smart_harness::decode_openai_response,
    )
    .await
    .unwrap_err();
    hang_guarded("the server observing the idle timeout", server)
        .await
        .unwrap();
    assert_eq!(
        crate::retry::classify(&error),
        crate::retry::Retryability::Retry,
        "a successful response body timeout remains retryable: {error:#}"
    );
    assert_eq!(
        crate::agentic::observability::error_class(&error),
        Some(crate::agentic::observability::ErrorClass::Timeout)
    );
    assert!(
        error
            .to_string()
            .contains("request failed reading response"),
        "{error}"
    );
}

/// Grounds the retained-byte reader in the primary and cap transports' real
/// reqwest idle timeout. The peer sends a complete rejection event, then waits
/// for the client to close its unfinished HTTP body; no fixture sleep completes it.
async fn observed_context_error_survives_core_idle_timeout(cap_summary: bool) {
    let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
    let uri = format!("http://{}", listener.local_addr().unwrap());
    let requests = Arc::new(Mutex::new(Vec::new()));
    let captured = requests.clone();
    let (closed, closure) = tokio::sync::oneshot::channel();
    let server = tokio::spawn(async move {
        let mut closed = Some(closed);
        loop {
            let (mut socket, _) = listener.accept().await.unwrap();
            let (path, body) = read_request(&mut socket).await;
            if path != "/v1/chat/completions" {
                socket
                    .write_all(
                        b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n",
                    )
                    .await
                    .unwrap();
                socket.shutdown().await.unwrap();
                continue;
            }
            let first = {
                let mut requests = captured.lock().unwrap();
                requests.push(body);
                requests.len() == 1
            };
            if !first {
                // An erroneous unchanged retry must fail promptly, allowing the
                // test to inspect the actual send count instead of hanging.
                socket
                    .write_all(b"HTTP/1.1 400 Bad Request\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")
                    .await
                    .unwrap();
                socket.shutdown().await.unwrap();
                return;
            }
            socket
                .write_all(b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: 4096\r\nConnection: close\r\n\r\ndata: {\"error\":{\"message\":\"Context size has been exceeded.\",\"code\":500}}\n\n")
                .await
                .unwrap();
            match socket.read(&mut [0_u8; 1]).await {
                Ok(0) => (),
                Err(error)
                    if matches!(
                        error.kind(),
                        std::io::ErrorKind::ConnectionReset | std::io::ErrorKind::ConnectionAborted
                    ) => {}
                other => panic!("the generation socket remained active: {other:?}"),
            }
            closed.take().unwrap().send(()).unwrap();
        }
    });
    let messages = msgs();
    let caveats = Caveats::top();
    let mut observation = observability::SolveObservation::default();
    let mut state = CompressState::new();
    let mut reason = None;
    let mut context = ctx(&uri, &messages, &caveats);
    context.action_nudges = false;
    context.inference_timeout_secs = 1;
    context.solve_obs = Some(&mut observation);
    context.compress_state = Some(&mut state);
    context.end_reason = Some(&mut reason);
    if cap_summary {
        context.max_tool_rounds = 0;
    }
    let result = hang_guarded(
        "the retained rejection ending the phase",
        chat_complete(context, &mut NoMcp),
    )
    .await;
    let socket_closed = hang_guarded("the idle timeout closing the socket", closure).await;
    server.abort();
    if let Err(error) = server.await {
        assert!(error.is_cancelled(), "fixture server failed: {error}");
    }
    socket_closed.expect("the peer must observe the close");
    if cap_summary {
        let (reply, _, _, _) = result.expect("a failed optional summary retains the cap fallback");
        assert!(reply.contains("tool-round limit (0"), "{reply}");
        assert_eq!(reason, Some(crate::TurnEndReason::RoundCap));
    } else {
        let error = result.expect_err("the protected operator prompt cannot shrink further");
        assert_eq!(
            crate::retry::classify(&error),
            crate::retry::Retryability::ContextExceeded,
            "the body idle timeout must not replace observed rejection evidence: {error:#}"
        );
    }
    let requests = requests.lock().unwrap();
    assert_eq!(
        requests.len(),
        1,
        "a timed-out rejection must not be resent unchanged"
    );
    assert_eq!(requests[0]["stream"], true);
    let rejections = observation
        .behavior_signals
        .iter()
        .filter(|signal| {
            matches!(
                signal,
                observability::BehaviorSignal::ContextExceeded { .. }
            )
        })
        .collect::<Vec<_>>();
    assert_eq!(
        rejections.len(),
        1,
        "the observed rejection must be recorded"
    );
    assert!(matches!(
        rejections[0],
        observability::BehaviorSignal::ContextExceeded {
            projected_tokens: None,
            ..
        }
    ));
    assert!(state.calibration.ratio(None) >= 1.5);
}

#[tokio::test]
async fn primary_stream_context_rejection_survives_the_body_idle_timeout() {
    observed_context_error_survives_core_idle_timeout(false).await;
}

#[tokio::test]
async fn cap_summary_stream_context_rejection_survives_the_body_idle_timeout() {
    observed_context_error_survives_core_idle_timeout(true).await;
}

fn has_tool_result(body: &Value) -> bool {
    body["messages"]
        .as_array()
        .unwrap()
        .iter()
        .any(|message| message["role"] == "tool")
}

fn identified(id: &str, mut frame: Value) -> Value {
    frame["id"] = json!(id);
    frame
}

/// A complete, valid tool batch except that its response ID changes mid-stream.
fn mixed_tool_batch() -> ResponseTemplate {
    wire(
        &[
            json!({"id":"resp-mixed-one","choices":[{"index":0,"delta":{"tool_calls":[{
                "index":0,"id":"call-a","type":"function","function":{"name":TOOL,"arguments":"{\"stage\":"}
            }]}}]}),
            json!({"id":"resp-mixed-two","choices":[{"index":0,"delta":{"tool_calls":[{
                "index":0,"function":{"arguments":"\"a\"}"}
            }]},"finish_reason":"tool_calls"}]}),
        ],
        true,
    )
}

fn mixed_answer() -> ResponseTemplate {
    wire(
        &[
            json!({"id":"resp-mixed-one","choices":[{"delta":{"content":"Fixture "}}]}),
            json!({"id":"resp-mixed-two","choices":[{"delta":{"content":"response."},"finish_reason":"stop"}]}),
        ],
        true,
    )
}

/// Mounts a generation mock and returns every request body it saw, in order.
async fn mount_rounds(
    server: &MockServer,
    respond: impl Fn(&Value, usize) -> ResponseTemplate + Send + Sync + 'static,
) -> Arc<Mutex<Vec<Value>>> {
    let seen = Arc::new(Mutex::new(Vec::<Value>::new()));
    let requests = seen.clone();
    Mock::given(method("POST"))
        .and(path("/v1/chat/completions"))
        .respond_with(move |request: &Request| {
            let body = body_json(request);
            let mut requests = requests.lock().unwrap();
            let same_kind = requests
                .iter()
                .filter(|seen| has_tool_result(seen) == has_tool_result(&body))
                .count();
            let response = respond(&body, same_kind);
            requests.push(body);
            response
        })
        .mount(server)
        .await;
    seen
}

async fn run_fixture_turn(
    uri: &str,
    tools: &mut FixtureTools,
    reason: &mut Option<crate::TurnEndReason>,
) -> anyhow::Result<(String, bool, Option<crate::TokenUsage>, u32)> {
    let messages = msgs();
    let caveats = Caveats::top();
    let mut context = ctx(uri, &messages, &caveats);
    context.action_nudges = false;
    let mut permission = FixtureMcpPermission::new(TOOL, &caveats);
    context.permission_gate = Some(&mut permission);
    context.end_reason = Some(reason);
    chat_complete(context, tools).await
}

/// #2334: a mixed response is rejected before its tool call can run, retried
/// once, and the clean second response runs the tool exactly once. The trap:
/// if the mixed batch had dispatched, the tool ledger would read `["a", "a"]`.
#[tokio::test]
async fn primary_stream_retries_a_mixed_response_id_once_before_its_tool_runs() {
    let server = MockServer::start().await;
    let seen = mount_rounds(&server, |body, nth| match (has_tool_result(body), nth) {
        (false, 0) => mixed_tool_batch(),
        (false, _) => wire(
            &[identified("resp-clean", complete_single_tool_frame())],
            true,
        ),
        (true, _) => streamed_answer(),
    })
    .await;
    let mut tools = FixtureTools::default();
    let mut reason = None;
    let (text, _, _, _) = run_fixture_turn(&server.uri(), &mut tools, &mut reason)
        .await
        .expect("one fresh attempt recovers a mixed response");
    assert_eq!(
        crate::agentic::model_explanation(&text),
        "Fixture response."
    );
    assert_eq!(tools.0, ["a"], "the mixed batch must never dispatch");
    let seen = seen.lock().unwrap();
    assert_eq!(
        seen.iter().filter(|body| !has_tool_result(body)).count(),
        2,
        "the tool round is sent exactly twice: rejected, then accepted"
    );
}

/// #2334: a tool that already ran in an earlier round is never replayed when a
/// later round's response mixes IDs; only that round's request is resent. The
/// accepted retry also carries a different valid ID from the first round.
#[tokio::test]
async fn a_mixed_response_after_a_dispatched_tool_retries_the_round_without_replay() {
    let server = MockServer::start().await;
    let seen = mount_rounds(&server, |body, nth| match (has_tool_result(body), nth) {
        (false, _) => wire(
            &[identified("resp-first-round", complete_single_tool_frame())],
            true,
        ),
        (true, 0) => mixed_answer(),
        (true, _) => wire(
            &[identified(
                "resp-second-round",
                json!({"choices":[{"delta":{"content":"Fixture response."},"finish_reason":"stop"}]}),
            )],
            true,
        ),
    })
    .await;
    let mut tools = FixtureTools::default();
    let mut reason = None;
    let (text, _, _, _) = run_fixture_turn(&server.uri(), &mut tools, &mut reason)
        .await
        .expect("the answer round recovers with one fresh attempt");
    assert_eq!(
        crate::agentic::model_explanation(&text),
        "Fixture response."
    );
    assert_eq!(tools.0, ["a"], "a dispatched tool is never replayed");
    let seen = seen.lock().unwrap();
    assert_eq!(seen.len(), 3, "tool round, mixed answer, one retry");
}

/// #2334: when the one fresh attempt is also mixed, the turn fails explicitly
/// after exactly two sends of that round, never as a completed turn, and the
/// earlier tool still ran exactly once.
#[tokio::test]
async fn a_mixed_response_on_the_retry_fails_after_exactly_two_attempts() {
    let server = MockServer::start().await;
    let seen = mount_rounds(&server, |body, _| {
        if has_tool_result(body) {
            mixed_answer()
        } else {
            wire(
                &[identified("resp-first-round", complete_single_tool_frame())],
                true,
            )
        }
    })
    .await;
    let mut tools = FixtureTools::default();
    let mut reason = None;
    let error = run_fixture_turn(&server.uri(), &mut tools, &mut reason)
        .await
        .expect_err("exhausted recovery is a failed turn");
    let text = format!("{error:#}");
    assert!(
        text.contains("stream changed response ID at data frame 2"),
        "{text}"
    );
    assert!(
        !text.contains("resp-mixed"),
        "IDs never reach the error: {text}"
    );
    assert_ne!(reason, Some(crate::TurnEndReason::Completed));
    assert_eq!(tools.0, ["a"]);
    let seen = seen.lock().unwrap();
    assert_eq!(
        seen.iter().filter(|body| has_tool_result(body)).count(),
        2,
        "the answer round is attempted exactly twice"
    );
}

#[path = "generation_bound_tests.rs"]
mod generation_bound_tests;

#[path = "provider_output_limit_tests.rs"]
mod provider_output_limit_tests;