haematite-cli 0.21.1

use std::io::{self, Cursor, Write};
use std::num::NonZeroU32;

use serde_json::value::RawValue;

use super::{Options, RelayError, Target, client_refusal, compact_line, read_line_bounded, relay};
use crate::client;
use crate::mcp::lys_pass::PassSource;
use crate::mcp::{adapter, rpc};
use crate::serve::protocol;

type TestResult = Result<(), Box<dyn std::error::Error>>;

const INITIALIZE: &str = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-11-25","capabilities":{},"clientInfo":{"name":"test","version":"1"}}}"#;
const QUERY: &str = r#"{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"run_query","arguments":{"selection":{"branch":"main","version":null},"query":".","page":{"rows":1,"bytes":4096},"after":null}}}"#;
const ANSWER: &str = r#"{"selection":{"branch":"main","version":{"generation":"1","sequence":"2"}},"rows":[{"value":1e2,"origin":null}],"next":null}"#;
const CAP: NonZeroU32 = NonZeroU32::MIN.saturating_add(1_048_575);

struct RecordingSender {
    requests: usize,
    ended: Option<client::Ended>,
}

impl adapter::Sender for RecordingSender {
    fn send(
        &mut self,
        request: protocol::Operation,
    ) -> Result<Box<RawValue>, adapter::ServiceRefusal> {
        assert!(matches!(request, protocol::Operation::QueryRun(_)));
        self.requests += 1;
        if let Some(ended) = self.ended.take() {
            return Err(client_refusal(ended));
        }
        serde_json::from_str(ANSWER).map_err(|error| adapter::ServiceRefusal {
            code: "test_answer_invalid".into(),
            message: error.to_string(),
            receipt: None,
        })
    }
}

fn sender() -> RecordingSender {
    RecordingSender {
        requests: 0,
        ended: None,
    }
}

fn lines(bytes: &[u8]) -> Result<Vec<serde_json::Value>, Box<dyn std::error::Error>> {
    let text = std::str::from_utf8(bytes)?;
    assert!(text.is_empty() || text.ends_with('\n'));
    text.lines()
        .map(|line| {
            let value: serde_json::Value = serde_json::from_str(line)?;
            assert_eq!(value["jsonrpc"], "2.0");
            assert!(value.get("id").is_some());
            assert!(value.get("result").is_some() || value.get("error").is_some());
            Ok(value)
        })
        .collect()
}

fn refusal(
    response: &serde_json::Value,
) -> Result<adapter::ServiceRefusal, Box<dyn std::error::Error>> {
    assert!(response["result"].get("structuredContent").is_some());
    let text = response["result"]["content"][0]["text"]
        .as_str()
        .ok_or("refusal text missing")?;
    Ok(serde_json::from_str(text)?)
}

#[test]
fn initialize_and_tools_list_answer_eleven_tools() -> TestResult {
    let input =
        format!("{INITIALIZE}\n{{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"tools/list\"}}\n");
    let mut output = Vec::new();
    let mut service = sender();
    relay(
        Cursor::new(input.as_bytes()),
        &mut output,
        &mut service,
        CAP,
    )?;
    let responses = lines(&output)?;
    assert_eq!(responses.len(), 2);
    assert_eq!(
        responses[0]["result"]["protocolVersion"],
        rpc::SUPPORTED_PROTOCOL_VERSION
    );
    assert_eq!(
        responses[1]["result"]["tools"].as_array().map(Vec::len),
        Some(11)
    );
    assert_eq!(service.requests, 0);
    Ok(())
}

#[test]
fn two_requests_produce_two_ordered_lines() -> TestResult {
    let mut output = Vec::new();
    relay(Cursor::new(b"{\"jsonrpc\":\"2.0\",\"id\":17,\"method\":\"ping\"}\n{\"jsonrpc\":\"2.0\",\"id\":18,\"method\":\"ping\"}\n"), &mut output, &mut sender(), CAP)?;
    let responses = lines(&output)?;
    assert_eq!(responses.len(), 2);
    assert_eq!(responses[0]["id"], 17);
    assert_eq!(responses[1]["id"], 18);
    Ok(())
}

#[test]
fn oversized_line_is_refused_and_next_line_is_answered() -> TestResult {
    let input = format!(
        "{}\n{{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"ping\"}}\n",
        "x".repeat(4096)
    );
    let mut reader = Cursor::new(input.as_bytes());
    let first = read_line_bounded(&mut reader, 64)?
        .ok_or_else(|| io::Error::other("expected oversized line"))?;
    assert_eq!(first.len(), 65);
    assert_eq!(
        read_line_bounded(&mut reader, 64)?.as_deref(),
        Some(b"{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"ping\"}".as_slice())
    );
    let mut output = Vec::new();
    relay(
        Cursor::new(input.as_bytes()),
        &mut output,
        &mut sender(),
        NonZeroU32::MIN.saturating_add(63),
    )?;
    let responses = lines(&output)?;
    assert_eq!(responses.len(), 2);
    assert_eq!(responses[0]["error"]["data"]["code"], "rpc_frame_too_large");
    assert_eq!(responses[1]["id"], 2);
    Ok(())
}

#[test]
fn end_of_input_returns_without_extra_output() -> TestResult {
    let mut output = Vec::new();
    relay(Cursor::new(b""), &mut output, &mut sender(), CAP)?;
    assert!(output.is_empty());
    relay(
        Cursor::new(b"{\"jsonrpc\":\"2.0\",\"id\":9,\"method\":\"ping\"}"),
        &mut output,
        &mut sender(),
        CAP,
    )?;
    assert_eq!(lines(&output)?.len(), 1);
    Ok(())
}

#[test]
fn notifications_write_nothing_and_send_nothing() -> TestResult {
    let input = b"{\"jsonrpc\":\"2.0\",\"method\":\"notifications/initialized\"}\n{\"jsonrpc\":\"2.0\",\"method\":\"ping\"}\n";
    let mut output = Vec::new();
    let mut service = sender();
    relay(Cursor::new(input), &mut output, &mut service, CAP)?;
    assert!(output.is_empty());
    assert_eq!(service.requests, 0);
    Ok(())
}

#[test]
fn cannot_connect_is_a_named_tool_result() -> TestResult {
    let mut service = RecordingSender {
        requests: 0,
        ended: Some(client::Ended::NotAsked("could not connect".into())),
    };
    let mut output = Vec::new();
    relay(
        Cursor::new(format!("{QUERY}\n")),
        &mut output,
        &mut service,
        CAP,
    )?;
    let responses = lines(&output)?;
    assert_eq!(responses.len(), 1);
    assert_eq!(responses[0]["result"]["isError"], true);
    assert_eq!(refusal(&responses[0])?.code, "mcp_client_not_asked");
    assert!(responses[0].get("error").is_none());
    assert_eq!(service.requests, 1);
    Ok(())
}

#[test]
fn unknown_outcome_has_its_own_code_and_no_retry() -> TestResult {
    let mut service = RecordingSender {
        requests: 0,
        ended: Some(client::Ended::Unknown("answer lost".into())),
    };
    let mut output = Vec::new();
    relay(
        Cursor::new(format!("{QUERY}\n")),
        &mut output,
        &mut service,
        CAP,
    )?;
    let responses = lines(&output)?;
    assert_eq!(refusal(&responses[0])?.code, "mcp_client_outcome_unknown");
    assert_eq!(responses[0]["result"]["isError"], true);
    assert_eq!(service.requests, 1);
    Ok(())
}

#[test]
fn output_contains_only_protocol_lines() -> TestResult {
    let input = format!("{INITIALIZE}\n{QUERY}\n[]\nnot-json\n");
    let mut output = Vec::new();
    relay(Cursor::new(input), &mut output, &mut sender(), CAP)?;
    assert_eq!(lines(&output)?.len(), 4);
    assert!(std::str::from_utf8(&output)?.contains("1e2"));
    Ok(())
}

#[test]
fn output_write_error_stops_before_another_request() -> TestResult {
    struct BrokenWriter {
        attempts: usize,
    }
    impl Write for BrokenWriter {
        fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
            assert!(!bytes.is_empty());
            self.attempts += 1;
            Err(io::Error::new(io::ErrorKind::BrokenPipe, "closed"))
        }
        fn flush(&mut self) -> io::Result<()> {
            Ok(())
        }
    }
    let mut writer = BrokenWriter { attempts: 0 };
    let mut service = sender();
    let outcome = relay(
        Cursor::new(format!("{QUERY}\n{QUERY}\n")),
        &mut writer,
        &mut service,
        CAP,
    );
    // Since the typed receipt work (BUSYSHARDNODEADLINE-001 Amendment 19:
    // MCP consumers retain the receipt) an output failure keeps the whole
    // response it could not write, so a mutation's answer is never lost.
    let (cause, response) = match outcome {
        Err(RelayError::ReceiptOutput { cause, response }) => (cause, response),
        other => return Err(format!("the output failure kept no response: {other:?}").into()),
    };
    assert_eq!(cause.kind(), io::ErrorKind::BrokenPipe);
    assert_eq!(response.last(), Some(&b'\n'));
    serde_json::from_slice::<serde_json::Value>(&response)?;
    assert_eq!(writer.attempts, 1);
    assert_eq!(service.requests, 1);
    Ok(())
}

#[test]
fn multiline_row_compacts_to_one_line_without_changing_tokens() -> TestResult {
    let mut frame =
        b"{\n \"rows\": [ {\"value\": { \"amount\": 1e2, \"id\": 123456789012345678 }} ]\r\n}"
            .to_vec();
    compact_line(&mut frame);
    assert_eq!(
        frame,
        br#"{"rows":[{"value":{"amount":1e2,"id":123456789012345678}}]}"#
    );
    assert!(!frame.contains(&b'\n'));
    serde_json::from_slice::<serde_json::Value>(&frame)?;
    Ok(())
}

#[test]
fn escaped_newline_and_spaces_inside_strings_keep_their_bytes() -> TestResult {
    let mut frame = b"{\n \"value\": \"a\\nb  \\t \\\"quoted\\\" \\\\ tail\"\n}".to_vec();
    compact_line(&mut frame);
    assert_eq!(frame, br#"{"value":"a\nb  \t \"quoted\" \\ tail"}"#);
    serde_json::from_slice::<serde_json::Value>(&frame)?;
    Ok(())
}

#[test]
fn multiline_service_row_keeps_text_and_structured_content_identical() -> TestResult {
    struct PrettySender;
    impl adapter::Sender for PrettySender {
        fn send(
            &mut self,
            request: protocol::Operation,
        ) -> Result<Box<RawValue>, adapter::ServiceRefusal> {
            assert!(matches!(request, protocol::Operation::QueryRun(_)));
            serde_json::from_str("{\n\"selection\":{\"branch\":\"main\",\"version\":{\"generation\":\"1\",\"sequence\":\"2\"}},\n\"rows\":[{\"value\":{\n\"amount\":1e2,\"id\":123456789012345678,\"text\":\"a\\nb  \\t tail\"\n},\"origin\":null}],\"next\":null\n}")
                .map_err(|error| adapter::ServiceRefusal { code: "test_answer_invalid".into(), message: error.to_string() , receipt: None})
        }
    }
    #[derive(serde::Deserialize)]
    struct Content {
        text: String,
    }
    #[derive(serde::Deserialize)]
    struct ResultBody {
        #[serde(rename = "structuredContent")]
        structured: Box<RawValue>,
        content: [Content; 1],
    }
    #[derive(serde::Deserialize)]
    struct Response {
        result: ResultBody,
    }
    let mut output = Vec::new();
    relay(
        Cursor::new(format!("{QUERY}\n")),
        &mut output,
        &mut PrettySender,
        CAP,
    )?;
    assert_eq!(lines(&output)?.len(), 1);
    let response: Response = serde_json::from_slice(&output)?;
    assert_eq!(
        response.result.content[0].text,
        response.result.structured.get()
    );
    assert!(response.result.structured.get().contains("1e2"));
    assert!(
        response
            .result
            .structured
            .get()
            .contains("123456789012345678")
    );
    assert!(response.result.structured.get().contains(r"a\nb  \t tail"));
    Ok(())
}

fn words(list: &[&str]) -> Vec<String> {
    list.iter().map(|word| (*word).to_owned()).collect()
}

#[test]
fn a_pass_file_is_taken_once_and_needs_a_path() -> TestResult {
    let options = Options::parse(words(&["--socket", "/s", "--pass-file", "/p"]).into_iter())?;
    assert!(matches!(options.target, Target::Socket(_)));
    assert!(
        matches!(&options.pass, Some(PassSource::File(path)) if path == std::path::Path::new("/p"))
    );
    let without = Options::parse(words(&["--socket", "/s"]).into_iter())?;
    assert!(without.pass.is_none());
    for refused in [
        &["--socket", "/s", "--pass-file"][..],
        &["--socket", "/s", "--pass-file", "--timeout-ms"][..],
        &["--socket", "/s", "--pass-file", "/p", "--pass-file", "/q"][..],
    ] {
        assert!(
            matches!(
                Options::parse(words(refused).into_iter()),
                Err(RelayError::Arguments(_))
            ),
            "{refused:?}"
        );
    }
    Ok(())
}

#[test]
fn a_lys_agent_takes_an_issuer_an_agent_and_an_absolute_credential_file() -> TestResult {
    let options = Options::parse(
        words(&[
            "--socket",
            "/s",
            "--lys-agent",
            "https://lys.example/",
            "agent-1",
            "--credential-file",
            "/c",
        ])
        .into_iter(),
    )?;
    let Some(PassSource::Lys(agent)) = &options.pass else {
        return Err("a lys agent was not taken as the pass source".into());
    };
    assert_eq!(agent.route(), "https://lys.example/api/agents/agent-1/pass");
    assert_eq!(agent.credential_file(), Some(std::path::Path::new("/c")));
    let local = Options::parse(
        words(&[
            "--credential-file",
            "/c",
            "--lys-agent",
            "http://127.0.0.1:8080",
            "a",
            "--config",
            "/f",
        ])
        .into_iter(),
    )?;
    assert!(matches!(local.pass, Some(PassSource::Lys(_))));
    Ok(())
}

#[test]
fn a_lys_agent_is_refused_without_its_parts_or_beside_a_pass_file() {
    let socket: &[&str] = &["--socket", "/s"];
    let pass_file: &[&str] = &["--pass-file", "/p"];
    let lys: &[&str] = &["--lys-agent", "https://lys.example", "agent-1"];
    let credential: &[&str] = &["--credential-file", "/c"];
    let with_issuer = |issuer: &'static str, agent: &'static str| {
        [socket, &["--lys-agent", issuer, agent][..], credential].concat()
    };
    let refused: Vec<Vec<&str>> = vec![
        // Exclusive with --pass-file, whichever order.
        [socket, pass_file, lys, credential].concat(),
        [socket, lys, credential, pass_file].concat(),
        [socket, pass_file, credential].concat(),
        // Both parts are required.
        [socket, lys].concat(),
        [socket, credential].concat(),
        [socket, &["--lys-agent", "https://lys.example"][..]].concat(),
        [
            socket,
            &["--lys-agent", "https://lys.example"][..],
            credential,
        ]
        .concat(),
        [socket, lys, &["--credential-file"][..]].concat(),
        // The credential path is absolute.
        [socket, lys, &["--credential-file", "c"][..]].concat(),
        // Taken once each.
        [socket, lys, lys, credential].concat(),
        [socket, lys, credential, credential].concat(),
        // The issuer is an https origin, or http only on 127.0.0.1 or localhost.
        with_issuer("http://lys.example", "a"),
        with_issuer("https://lys.example/x", "a"),
        with_issuer("https://lys.example?x=1", "a"),
        with_issuer("https://u@lys.example", "a"),
        with_issuer("lys.example", "a"),
        // The agent id is plain path text.
        with_issuer("https://lys.example", "a/b"),
    ];
    for refused in refused {
        assert!(
            matches!(
                Options::parse(words(&refused).into_iter()),
                Err(RelayError::Arguments(_))
            ),
            "{refused:?}"
        );
    }
}

#[test]
fn a_lys_agent_takes_a_run_pass_variable_in_place_of_a_file() -> TestResult {
    let options = Options::parse(
        words(&[
            "--socket",
            "/s",
            "--lys-agent",
            "https://lys.example/",
            "agent-1",
            "--run-pass-env",
            "LYS_AGENT_PASS",
        ])
        .into_iter(),
    )?;
    let Some(PassSource::Lys(agent)) = &options.pass else {
        return Err("a lys agent was not taken as the pass source".into());
    };
    assert_eq!(agent.credential_file(), None);
    Ok(())
}