link-assistant-router 1.12.0

use super::*;

/// Translation of an `OpenAI` chat body into Gemini's request shape,
/// including the system-instruction split and generation config.
#[test]
fn chat_requests_translate_into_the_gemini_shape() {
    let body = json!({
        "messages": [
            {"role": "system", "content": "be brief"},
            {"role": "user", "content": "hello"},
            {"role": "assistant", "content": "hi", "tool_calls": [{
                "id": "call_1", "type": "function",
                "function": {"name": "lookup", "arguments": "{}"}
            }]},
            {"role": "tool", "tool_call_id": "call_1", "content": "result"}
        ],
        "max_tokens": 128,
        "temperature": 0.4,
        "top_p": 0.9
    });
    let request = chat_to_gemini_request(&body);

    assert_eq!(request["systemInstruction"]["parts"][0]["text"], "be brief");
    let contents = request["contents"].as_array().expect("contents");
    assert_eq!(contents.len(), 3, "system is lifted out of the turn list");
    assert_eq!(contents[0]["role"], "user");
    assert_eq!(contents[1]["role"], "model", "assistant maps to model");
    assert_eq!(
        contents[2]["role"], "user",
        "tool results map to a user turn"
    );
    assert_eq!(request["generationConfig"]["maxOutputTokens"], 128);
    assert_eq!(request["generationConfig"]["temperature"], 0.4);
    assert_eq!(request["generationConfig"]["topP"], 0.9);
}

#[test]
fn a_request_without_knobs_omits_the_generation_config() {
    let request = chat_to_gemini_request(&json!({"messages": []}));
    assert!(request.get("generationConfig").is_none());
    assert!(request.get("systemInstruction").is_none());
}

#[test]
fn max_completion_tokens_is_accepted_as_the_output_cap() {
    let request = chat_to_gemini_request(&json!({
        "messages": [],
        "max_completion_tokens": 64
    }));
    assert_eq!(request["generationConfig"]["maxOutputTokens"], 64);
}

#[test]
fn the_code_assist_envelope_carries_the_model() {
    let envelope = code_assist_envelope("models/nimbus-3-flash", &json!({"contents": []}));
    assert_eq!(envelope["model"], "nimbus-3-flash");
    assert_eq!(envelope["request"]["contents"], json!([]));
    let only_one = code_assist_envelope("models/models/exact", &json!({"contents": []}));
    assert_eq!(only_one["model"], "models/exact");
}

/// Responses are translated back into the `OpenAI` completion shape, with
/// usage carried across so spend accounting stays truthful.
#[test]
fn gemini_responses_translate_back_with_usage() {
    let response = json!({
        "modelVersion": "nimbus-3-flash-actual",
        "candidates": [{
            "content": {"parts": [{"text": "one "}, {"text": "two"}]},
            "finishReason": "STOP"
        }],
        "usageMetadata": {"promptTokenCount": 11, "candidatesTokenCount": 5}
    });
    let chat = gemini_response_to_chat(&response, "nimbus-3-flash");

    assert_eq!(chat["model"], "nimbus-3-flash-actual");
    assert_eq!(chat["choices"][0]["message"]["content"], "one two");
    assert_eq!(chat["choices"][0]["finish_reason"], "stop");
    assert_eq!(chat["usage"]["prompt_tokens"], 11);
    assert_eq!(chat["usage"]["completion_tokens"], 5);
    assert_eq!(chat["usage"]["total_tokens"], 16);
}

/// Code Assist nests the payload under `response`; both shapes are read.
#[test]
fn a_nested_code_assist_response_is_unwrapped() {
    let nested = json!({
        "response": {
            "candidates": [{"content": {"parts": [{"text": "inner"}]}}]
        }
    });
    let chat = gemini_response_to_chat(&nested, "nimbus-3-flash");
    assert_eq!(chat["choices"][0]["message"]["content"], "inner");
}

#[test]
fn finish_reasons_map_onto_the_openai_vocabulary() {
    assert_eq!(map_finish_reason("MAX_TOKENS"), "length");
    for blocked in ["SAFETY", "RECITATION", "BLOCKLIST", "PROHIBITED_CONTENT"] {
        assert_eq!(map_finish_reason(blocked), "content_filter", "{blocked}");
    }
    assert_eq!(map_finish_reason("STOP"), "stop");
    assert_eq!(map_finish_reason("SOMETHING_NEW"), "content_filter");
}

#[test]
fn buffered_tool_calls_preserve_identity_arguments_finish_and_usage_on_both_surfaces() {
    let gemini = json!({
        "candidates": [{
            "content": {"parts": [{"functionCall": {
                "id": "call_7", "name": "lookup", "args": {"key": "value"}
            }}]},
            "finishReason": "STOP"
        }],
        "usageMetadata": {"promptTokenCount": 2, "candidatesTokenCount": 3}
    });
    let chat = gemini_response_to_chat(&gemini, "served-model");
    assert_eq!(chat["choices"][0]["finish_reason"], "tool_calls");
    let call = &chat["choices"][0]["message"]["tool_calls"][0];
    assert_eq!(call["id"], "call_7");
    assert_eq!(call["function"]["name"], "lookup");
    assert_eq!(call["function"]["arguments"], "{\"key\":\"value\"}");
    assert_eq!(chat["usage"]["total_tokens"], 5);

    let response = responses::from_chat(
        &chat,
        "requested-model",
        responses::Finish::from_gemini(gemini_finish_reason(&gemini).unwrap()),
    );
    assert_eq!(response["status"], "completed");
    assert_eq!(response["output"][0]["type"], "function_call");
    assert_eq!(response["output"][0]["call_id"], "call_7");
    assert_eq!(response["output"][0]["name"], "lookup");
    assert_eq!(response["output"][0]["arguments"], "{\"key\":\"value\"}");
    assert_eq!(response["usage"]["total_tokens"], 5);
}

#[test]
fn buffered_prompt_blocks_do_not_become_successful_empty_responses() {
    let gemini = json!({"promptFeedback": {"blockReason": "SAFETY"}});
    let chat = gemini_response_to_chat(&gemini, "served-model");
    assert_eq!(chat["choices"][0]["finish_reason"], "content_filter");
    let response = responses::from_chat(
        &chat,
        "requested-model",
        responses::Finish::from_gemini(gemini_finish_reason(&gemini).unwrap()),
    );
    assert_eq!(response["status"], "incomplete");
    assert_eq!(
        response["incomplete_details"],
        json!({"reason": "content_filter"})
    );
}

#[test]
fn message_text_is_extracted_from_both_content_shapes() {
    assert_eq!(extract_message_text(Some(&json!("plain"))), "plain");
    assert_eq!(
        extract_message_text(Some(&json!([{"text": "a"}, {"text": "b"}]))),
        "ab"
    );
    assert_eq!(extract_message_text(Some(&json!(["a", "b"]))), "ab");
    assert_eq!(extract_message_text(None), "");
    assert_eq!(extract_message_text(Some(&json!(42))), "");
}

/// Incremental output reports the concrete upstream identity without
/// Router-private metadata.
#[test]
fn incremental_chat_stream_preserves_served_model_only() {
    let mut translator = stream::OpenAiStreamTranslator::new("catalog-alias");
    let payload = translator
        .push(
            br#"data: {"response":{"modelVersion":"future-upstream-model","candidates":[{"content":{"parts":[{"text":"hello"}]},"finishReason":"STOP"}]}}

"#,
        )
        .expect("translate Gemini SSE");
    let payload = String::from_utf8(payload.to_vec()).expect("UTF-8 SSE");
    let chunks: Vec<Value> = payload
        .lines()
        .filter_map(|line| line.strip_prefix("data: "))
        .filter(|data| *data != "[DONE]")
        .map(|data| serde_json::from_str(data).expect("JSON SSE frame"))
        .collect();
    assert_eq!(chunks.len(), 3);
    for chunk in chunks {
        assert_eq!(chunk["model"], "future-upstream-model");
        assert!(chunk.get("x_router_upstream_model").is_none());
    }
    assert!(payload.ends_with("data: [DONE]\n\n"));
}