newt-core 0.8.0

use super::*;

/// Records whether every inference request preserves normal multi-turn wire
/// ordering while also carrying the protected recovery copy near the system
/// prompt.
struct ActivePromptWireOrderResponder {
    exact_task: &'static str,
    requests: Arc<AtomicUsize>,
    order_valid: Arc<AtomicBool>,
}

impl Respond for ActivePromptWireOrderResponder {
    fn respond(&self, req: &Request) -> ResponseTemplate {
        let body = body_json(req);
        let messages = body["messages"].as_array().cloned().unwrap_or_default();
        let card_index = messages.iter().position(|message| {
            message["role"].as_str() == Some("system")
                && message["content"]
                    .as_str()
                    .is_some_and(|text| text.starts_with(prompt_read::ACTIVE_PROMPT_PREFIX))
        });
        let protected_copy_is_earlier = card_index.is_some_and(|index| {
            index + 1 < messages.len().saturating_sub(1)
                && messages[index + 1]["role"].as_str() == Some("user")
                && messages[index + 1]["content"].as_str() == Some(self.exact_task)
        });
        let live_task_is_newest = messages.last().is_some_and(|message| {
            message["role"].as_str() == Some("user")
                && message["content"].as_str() == Some(self.exact_task)
        });
        let exact_task_copies = messages
            .iter()
            .filter(|message| {
                message["role"].as_str() == Some("user")
                    && message["content"].as_str() == Some(self.exact_task)
            })
            .count();
        self.order_valid.fetch_and(
            protected_copy_is_earlier && live_task_is_newest && exact_task_copies == 2,
            Ordering::SeqCst,
        );
        self.requests.fetch_add(1, Ordering::SeqCst);

        if is_stream(req) {
            ndjson(&[serde_json::json!({
                "message": {"content": "wire order preserved"},
                "done": true,
                "prompt_eval_count": 20,
                "eval_count": 4
            })])
        } else {
            ResponseTemplate::new(200).set_body_json(serde_json::json!({
                "message": {"content": "wire order preserved"},
                "prompt_eval_count": 20,
                "eval_count": 4
            }))
        }
    }
}

#[tokio::test]
async fn multi_turn_wire_keeps_live_operator_task_newest_after_protected_copy() {
    const CURRENT_TASK: &str = "CURRENT OPERATOR TASK: report the exact wire ordering";
    let server = MockServer::start().await;
    let requests = Arc::new(AtomicUsize::new(0));
    let order_valid = Arc::new(AtomicBool::new(true));
    Mock::given(method("POST"))
        .and(path("/api/chat"))
        .respond_with(ActivePromptWireOrderResponder {
            exact_task: CURRENT_TASK,
            requests: requests.clone(),
            order_valid: order_valid.clone(),
        })
        .mount(&server)
        .await;

    let messages = vec![
        MemMessage::system("you are a test"),
        MemMessage::user("an older operator request"),
        MemMessage::assistant("the older request is complete"),
        MemMessage::user(CURRENT_TASK),
    ];
    let caveats = Caveats::top();
    let uri = server.uri();
    let mut c = ctx(&uri, &messages, &caveats);
    c.task = CURRENT_TASK;
    let (reply, streamed, _, _) = chat_complete(c, &mut NoMcp)
        .await
        .expect("ordinary multi-turn request should complete");

    assert_eq!(
        crate::agentic::model_explanation(&reply),
        "wire order preserved"
    );
    assert!(!streamed, "the host renders the accepted answer (#2372)");
    assert_eq!(
        requests.load(Ordering::SeqCst),
        1,
        "one generation; no display reissue (#2372)"
    );
    assert!(
        order_valid.load(Ordering::SeqCst),
        "every request must keep the protected copy earlier and the live task newest"
    );
}

// -----------------------------------------------------------------------
// OpenAI-path coverage
// -----------------------------------------------------------------------

/// Large MCP surface from the field transcript: 169 remote definitions before
/// Newt's built-ins. None is invoked; this fixture grounds the request-envelope
/// regression without a live MCP server.
struct ManyToolsMcp {
    count: usize,
}

#[async_trait::async_trait]
impl McpTools for ManyToolsMcp {
    fn handles(&self, _name: &str) -> bool {
        false
    }

    fn tool_defs(&self) -> Vec<serde_json::Value> {
        (0..self.count)
            .map(|index| {
                serde_json::json!({
                    "type": "function",
                    "function": {
                        "name": format!("field_mcp__tool_{index:03}"),
                        "description": "Synthetic remote tool for the provider-envelope regression.",
                        "parameters": {"type": "object", "properties": {}}
                    }
                })
            })
            .collect()
    }

    async fn call(&mut self, _leased: &LeasedMcpCall<'_>) -> String {
        "unexpected MCP call".to_string()
    }
}

/// Simulates an OpenAI-compatible gateway that rejects more than 128 function
/// tools with the same opaque invalid_argument shape seen in the transcript.
struct OpenAiToolEnvelopeResponder {
    max_seen: Arc<AtomicUsize>,
    kernel_seen: Arc<AtomicBool>,
}

impl Respond for OpenAiToolEnvelopeResponder {
    fn respond(&self, req: &Request) -> ResponseTemplate {
        let body = body_json(req);
        let tools = body["tools"].as_array().cloned().unwrap_or_default();
        self.max_seen.fetch_max(tools.len(), Ordering::SeqCst);
        let names = tools
            .iter()
            .filter_map(|tool| {
                tool.pointer("/function/name")
                    .and_then(serde_json::Value::as_str)
            })
            .collect::<Vec<_>>();
        self.kernel_seen.store(
            names.contains(&"run_command") && names.contains(&"tool_search"),
            Ordering::SeqCst,
        );
        if tools.len() > crate::agentic::tools::exposure::OPENAI_COMPATIBLE_MAX_FUNCTION_TOOLS {
            return ResponseTemplate::new(400).set_body_json(serde_json::json!({
                "error": {
                    "code": "invalid_argument",
                    "message": "an internal error occurred",
                    "type": "invalid_request_error"
                }
            }));
        }
        ResponseTemplate::new(200).set_body_json(serde_json::json!({
            "choices": [{"message": {"content": "tool envelope accepted"}}]
        }))
    }
}

#[tokio::test]
async fn openai_large_mcp_catalog_stays_within_wire_contract_and_keeps_shell() {
    let server = MockServer::start().await;
    let max_seen = Arc::new(AtomicUsize::new(0));
    let kernel_seen = Arc::new(AtomicBool::new(false));
    Mock::given(method("POST"))
        .and(path("/v1/chat/completions"))
        .respond_with(OpenAiToolEnvelopeResponder {
            max_seen: max_seen.clone(),
            kernel_seen: kernel_seen.clone(),
        })
        .mount(&server)
        .await;

    let messages = msgs();
    let caveats = Caveats::top();
    let uri = server.uri();
    let mut c = ctx(&uri, &messages, &caveats);
    c.kind = BackendKind::Openai;
    c.task = "test \"gh auth status\" now";
    let mut mcp = ManyToolsMcp { count: 169 };

    let (reply, _, _, _) = chat_complete(c, &mut mcp)
        .await
        .expect("the provider-compatible request should reach inference");

    assert_eq!(
        crate::agentic::model_explanation(&reply),
        "tool envelope accepted"
    );
    assert_eq!(
        max_seen.load(Ordering::SeqCst),
        crate::agentic::tools::exposure::OPENAI_COMPATIBLE_MAX_FUNCTION_TOOLS
    );
    assert!(
        kernel_seen.load(Ordering::SeqCst),
        "run_command and tool_search must survive optional MCP clipping"
    );
}

#[tokio::test]
async fn chat_complete_dispatches_openai_kind_and_returns_first_round_answer() {
    let server = MockServer::start().await;
    Mock::given(method("POST"))
        .and(path("/v1/chat/completions"))
        .and(header("authorization", "Bearer sk-test"))
        .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
            "choices": [{"message": {"content": "openai says hi"}}],
            "usage": {"prompt_tokens": 10, "completion_tokens": 4},
        })))
        .mount(&server)
        .await;

    let messages = msgs();
    let caveats = Caveats::top();
    let uri = server.uri();
    let mut c = ctx(&uri, &messages, &caveats);
    c.kind = BackendKind::Openai;
    c.api_key = Some("sk-test");
    // Calling chat_complete (not openai_chat_complete) pins the dispatch.
    let (reply, streamed, usage, hallu) = chat_complete(c, &mut NoMcp)
        .await
        .expect("openai dispatch should succeed");

    assert_eq!(crate::agentic::model_explanation(&reply), "openai says hi");
    assert!(
        !streamed,
        "a JSON fallback is not a streamed display response"
    );
    let u = usage.unwrap();
    assert_eq!((u.input_tokens, u.output_tokens), (10, 4));
    assert_eq!(hallu, 0);
}

#[test]
fn openai_progress_label_exposes_attempt_and_idle_timeout() {
    assert_eq!(
        inference_progress_label("kimi-k3", 2, 2, 120),
        "waiting for kimi-k3 · attempt 2/2 · 120s idle timeout…"
    );
}

/// Mirrors vLLM chat templates that accept exactly one system message and
/// require it to be the first message in the request.
struct StrictVllmSystemResponder;

impl Respond for StrictVllmSystemResponder {
    fn respond(&self, req: &Request) -> ResponseTemplate {
        let body = body_json(req);
        let messages = body["messages"].as_array().cloned().unwrap_or_default();
        let system_messages: Vec<(usize, &serde_json::Value)> = messages
            .iter()
            .enumerate()
            .filter(|(_, message)| message["role"].as_str() == Some("system"))
            .collect();
        let valid_system = system_messages.len() == 1
            && system_messages[0].0 == 0
            && system_messages[0].1["content"]
                .as_str()
                .is_some_and(|content| {
                    content.contains("you are a test")
                        && content.contains(prompt_read::ACTIVE_PROMPT_PREFIX)
                });
        let expected_tail = [
            ("user", "hello"),
            ("user", "an earlier request"),
            ("assistant", "the earlier request is complete"),
            ("user", "hello"),
        ];
        let valid_tail = messages.get(1..).is_some_and(|tail| {
            tail.len() == expected_tail.len()
                && tail
                    .iter()
                    .zip(expected_tail)
                    .all(|(message, (role, content))| {
                        message["role"].as_str() == Some(role)
                            && message["content"].as_str() == Some(content)
                    })
        });

        if !valid_system || !valid_tail {
            return ResponseTemplate::new(400).set_body_json(serde_json::json!({
                "error": {
                    "message": "System message must be at the beginning.",
                    "type": "BadRequestError"
                }
            }));
        }

        ResponseTemplate::new(200).set_body_json(serde_json::json!({
            "choices": [{"message": {"content": "strict vllm accepted the request"}}]
        }))
    }
}

#[tokio::test]
async fn openai_chat_coalesces_system_cards_for_strict_vllm_templates() {
    let server = MockServer::start().await;
    Mock::given(method("POST"))
        .and(path("/v1/chat/completions"))
        .respond_with(StrictVllmSystemResponder)
        .mount(&server)
        .await;

    let messages = vec![
        MemMessage::system("you are a test"),
        MemMessage::user("an earlier request"),
        MemMessage::assistant("the earlier request is complete"),
        MemMessage::user("hello"),
    ];
    let caveats = Caveats::top();
    let uri = server.uri();
    let mut c = ctx(&uri, &messages, &caveats);
    c.kind = BackendKind::Openai;
    c.task = "hello";

    let (reply, _, _, _) = chat_complete(c, &mut NoMcp)
        .await
        .expect("strict vLLM chat template should accept the next message after a backend switch");

    assert_eq!(
        crate::agentic::model_explanation(&reply),
        "strict vllm accepted the request"
    );
}

/// Shared composition rendering must preserve the existing wire diagnostic.
#[test]
fn openai_chat_rejects_system_messages_after_conversation_history() {
    let messages = vec![
        serde_json::json!({"role": "system", "content": "base policy"}),
        serde_json::json!({"role": "user", "content": "hello"}),
        serde_json::json!({"role": "system", "content": "late policy"}),
    ];

    let error = openai_chat_wire_messages(&messages)
        .expect_err("a late system message must not be promoted across user history");

    assert_eq!(
        error.to_string(),
        "invalid OpenAI chat message order: system messages must precede conversation history"
    );
}

/// The shared renderer must not expose its internal proposal-error prefix.
#[test]
fn openai_chat_coalescing_preserves_nontext_system_diagnostic() {
    let messages = vec![
        serde_json::json!({"role": "system", "content": "base policy"}),
        serde_json::json!({"role": "system", "content": []}),
    ];
    assert_eq!(
        openai_chat_wire_messages(&messages)
            .unwrap_err()
            .to_string(),
        "invalid OpenAI chat system message: content must be text before coalescing"
    );
}

#[tokio::test]
async fn openai_chat_projects_cognition_only_for_an_explicitly_capable_endpoint() {
    let server = MockServer::start().await;
    let request = Arc::new(Mutex::new(None));
    Mock::given(method("POST"))
        .and(path("/v1/chat/completions"))
        .respond_with(CaptureOpenAiRequestResponder {
            request: request.clone(),
        })
        .mount(&server)
        .await;

    let messages = msgs();
    let caveats = Caveats::top();
    let uri = server.uri();
    let mut c = ctx(&uri, &messages, &caveats);
    c.kind = BackendKind::Openai;
    c.cognition = Some(crate::role_profile::Cognition::Thoughtful);
    c.chat_completions_capability = crate::model_card::ChatCompletionsCapability {
        cognition: Some(true),
        chat_template_kwargs: Some(true),
        parallel_tool_calls: Some(false),
        bounded_reasoning_continuation: Some(true),
    };

    chat_complete(c, &mut NoMcp)
        .await
        .expect("capable OpenAI-compatible dispatch succeeds");

    let request = request
        .lock()
        .expect("capture lock")
        .clone()
        .expect("request captured");
    assert_eq!(request["max_tokens"], 10_000);
    assert_eq!(request["temperature"], 0.6);
    assert_eq!(request["top_p"], 0.95);
    assert_eq!(request["parallel_tool_calls"], false);
    assert_eq!(request["chat_template_kwargs"]["enable_thinking"], true);
    assert_eq!(
        request["chat_template_kwargs"]["truncate_history_thinking"],
        true
    );
}

#[tokio::test]
async fn openai_chat_omits_local_cognition_fields_for_an_unknown_endpoint() {
    let server = MockServer::start().await;
    let request = Arc::new(Mutex::new(None));
    Mock::given(method("POST"))
        .and(path("/v1/chat/completions"))
        .respond_with(CaptureOpenAiRequestResponder {
            request: request.clone(),
        })
        .mount(&server)
        .await;

    let messages = msgs();
    let caveats = Caveats::top();
    let uri = server.uri();
    let mut c = ctx(&uri, &messages, &caveats);
    c.kind = BackendKind::Openai;
    c.cognition = Some(crate::role_profile::Cognition::Meticulous);
    // #2312 (A6): an explicit allowance is not a declaration that the endpoint
    // accepts cognition knobs. The independent generation limit still applies.
    c.output_allowance = Some(12_000);

    chat_complete(c, &mut NoMcp)
        .await
        .expect("strict-compatible dispatch succeeds");

    let request = request
        .lock()
        .expect("capture lock")
        .clone()
        .expect("request captured");
    assert_eq!(request["max_tokens"], 12_000);
    for field in [
        "temperature",
        "top_p",
        "parallel_tool_calls",
        "chat_template_kwargs",
    ] {
        assert!(
            request.get(field).is_none(),
            "unknown endpoints must not receive `{field}`"
        );
    }
}

/// Dispatch one capable Chat Completions turn at `cognition` with the given
/// explicit output allowance and return every request body the server saw.
async fn capable_chat_bodies(
    cognition: Option<crate::role_profile::Cognition>,
    output_allowance: Option<u32>,
) -> Vec<serde_json::Value> {
    let server = MockServer::start().await;
    Mock::given(method("POST"))
        .and(path("/v1/chat/completions"))
        .respond_with(CaptureOpenAiRequestResponder {
            request: Arc::new(Mutex::new(None)),
        })
        .mount(&server)
        .await;
    let messages = msgs();
    let caveats = Caveats::top();
    let uri = server.uri();
    let mut c = ctx(&uri, &messages, &caveats);
    c.kind = BackendKind::Openai;
    c.cognition = cognition;
    c.output_allowance = output_allowance;
    c.chat_completions_capability = crate::model_card::ChatCompletionsCapability {
        cognition: Some(true),
        chat_template_kwargs: Some(true),
        parallel_tool_calls: Some(false),
        bounded_reasoning_continuation: Some(true),
    };
    openai_chat_complete(c, &mut NoMcp)
        .await
        .expect("capable OpenAI-compatible dispatch succeeds");
    // The active-prompt card mints a fresh address per turn; pin it so body
    // equality compares what the allowance could change.
    let address = regex::Regex::new(r"prompt:[0-9a-f-]{36}").expect("regex");
    server
        .received_requests()
        .await
        .expect("wiremock request journal")
        .iter()
        .map(|request| {
            let body = String::from_utf8_lossy(&request.body);
            let mut body: serde_json::Value =
                serde_json::from_str(&address.replace_all(&body, "prompt:ID")).expect("json body");
            // Output capacity now changes its host planning statement too.
            // Verify that exact statement before normalizing it; keep every
            // other prompt byte and generation field in the equality check.
            let cap = u32::try_from(body["max_tokens"].as_u64().unwrap()).unwrap();
            let guidance = output_budget::guidance(
                Some(cap),
                None,
                std::path::Path::new(NO_CHECKS_WORKSPACE),
                &crate::Scope::All,
                crate::tokens::TokenEstimation::default(),
            );
            let mut found = 0;
            for message in body["messages"].as_array_mut().unwrap() {
                if let Some(text) = message["content"]
                    .as_str()
                    .filter(|text| text.contains(&guidance))
                {
                    assert_eq!(message["role"], "system");
                    found += text.matches(&guidance).count();
                    message["content"] =
                        serde_json::json!(text.replace(&guidance, "[verified output planning]"));
                }
            }
            assert_eq!(found, 1, "one accurate host planning note");
            body
        })
        .collect()
}

/// #2312 (A1/A2): output allowance changes the cap and its planning note. The
/// cognition dial, thinking enablement and sampling stay byte-identical, so an
/// operator can compare two allowances at one cognition level. With no
/// override the cognition table still supplies the historical 10,000.
#[tokio::test]
async fn openai_chat_output_allowance_varies_cap_and_planning_guidance() {
    let thoughtful = Some(crate::role_profile::Cognition::Thoughtful);
    let default_bodies = capable_chat_bodies(thoughtful, None).await;
    let small_bodies = capable_chat_bodies(thoughtful, Some(3_000)).await;
    let large_bodies = capable_chat_bodies(thoughtful, Some(12_000)).await;
    assert!(!default_bodies.is_empty());
    assert_eq!(default_bodies.len(), small_bodies.len());
    assert_eq!(default_bodies.len(), large_bodies.len());

    for ((default, small), large) in default_bodies
        .into_iter()
        .zip(small_bodies)
        .zip(large_bodies)
    {
        let [(default_cap, default), (small_cap, small), (large_cap, large)] =
            [default, small, large].map(|mut body| {
                let cap = body
                    .as_object_mut()
                    .expect("object body")
                    .remove("max_tokens");
                (cap, body)
            });
        assert_eq!(default_cap, Some(serde_json::json!(10_000)));
        assert_eq!(small_cap, Some(serde_json::json!(3_000)));
        assert_eq!(large_cap, Some(serde_json::json!(12_000)));
        // With cap and verified planning normalized, temperature, top_p, chat_template_kwargs
        // and every other field are identical across the three allowances.
        assert_eq!(default["temperature"], 0.6);
        assert_eq!(default["chat_template_kwargs"]["enable_thinking"], true);
        assert_eq!(small, default);
        assert_eq!(large, default);
    }
}

/// #2312 (A1, no dial): a cognition-projecting endpoint accepts the cap field,
/// so an explicit allowance is SENT even when no cognition dial is set, while
/// thinking and sampling stay unset. Without this arm the only way to put a
/// cap on the wire was to set a dial, which also moves thinking/temperature/
/// top_p. With no allowance the independent generation maximum still applies.
#[tokio::test]
async fn openai_chat_output_allowance_is_sent_without_a_cognition_dial() {
    let unset = capable_chat_bodies(None, None).await;
    let capped = capable_chat_bodies(None, Some(3_000)).await;
    assert!(!unset.is_empty());
    assert_eq!(unset.len(), capped.len());
    for (mut unset, mut capped) in unset.into_iter().zip(capped) {
        assert_eq!(
            unset.as_object_mut().unwrap().remove("max_tokens"),
            Some(serde_json::json!(16_384))
        );
        for field in ["temperature", "top_p", "chat_template_kwargs"] {
            assert!(
                unset.get(field).is_none(),
                "no allowance, no dial: `{field}`"
            );
        }
        assert_eq!(
            capped
                .as_object_mut()
                .expect("object body")
                .remove("max_tokens"),
            Some(serde_json::json!(3_000))
        );
        assert_eq!(
            capped, unset,
            "only max_tokens and its verified planning guidance may differ"
        );
    }
}

/// #2747: a neighboring turn's one-second budget must not change this catalog.
/// Repeats the actual wire assertion beside the real timeout verification fixture.
#[cfg(unix)]
#[test]
fn command_budget_2747_concurrent_wire_contract() {
    std::thread::scope(|scope| {
        scope
            .spawn(super::verification_outcomes::every_verification_case_ends_within_its_allowance);
        for _ in 0..30 {
            openai_chat_output_allowance_varies_cap_and_planning_guidance();
        }
    });
}

/// #2775: independent captured session dialects reach the actual provider wire;
/// catalog construction must not replace them with this host's environment.
#[tokio::test]
async fn shell_dialect_2775_captured_session_reaches_wire() {
    for windows_host in [true, false, true] {
        let server = MockServer::start().await;
        Mock::given(method("POST"))
            .and(path("/v1/chat/completions"))
            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
                "choices": [{"message": {"content": "ready"}}],
            })))
            .mount(&server)
            .await;
        let messages = msgs();
        let caveats = Caveats::top();
        let uri = server.uri();
        let mut c = ctx(&uri, &messages, &caveats);
        c.kind = BackendKind::Openai;
        c.shell_dialect = shell_dialect_sentence(true, windows_host);
        let expected = c.shell_dialect;
        chat_complete(c, &mut NoMcp).await.unwrap();
        let requests = server.received_requests().await.unwrap();
        assert!(!requests.is_empty());
        for request in requests {
            let body = body_json(&request);
            let tools = body["tools"].as_array().unwrap();
            let shell = tools
                .iter()
                .find(|t| t["function"]["name"] == "run_command")
                .unwrap();
            assert!(shell["function"]["description"]
                .as_str()
                .unwrap()
                .contains(expected));
        }
    }
}