haematite-cli 0.21.8

The haematite command: serve a store, its door and Studio, answer MCP, query with jq, and host the store's assistant seat
Documentation
use std::collections::BTreeSet;

use serde_json::{Value, json};

use super::{check, parser, words};
use crate::serve::authorize;
use crate::serve::protocol::Operation;
use crate::serve::route_matrix;

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

fn operation(method: &str, params: &Value) -> Result<Operation, serde_json::Error> {
    serde_json::from_str(&json!({"method": method, "params": params}).to_string())
}

fn answer(method: &str, params: &Value) -> Result<Value, Box<dyn std::error::Error>> {
    match operation(method, params)? {
        Operation::QueryCheck(request) => Ok(check::check(&request)),
        Operation::QueryWords(request) => Ok(check::words(&request)?),
        other => Err(format!("unexpected operation: {other:?}").into()),
    }
}

#[test]
fn query_check_gives_source_sort_and_limit_byte_spans() -> TestResult {
    assert_eq!(
        answer(
            "query.check",
            &json!({"query": "patients | sort(.a) | limit(2)"})
        )?,
        json!({
            "stages": [
                {"word": "patients", "span": {"start": 0, "end": 8}},
                {"word": "sort", "span": {"start": 11, "end": 19}},
                {"word": "limit", "span": {"start": 22, "end": 30}}
            ],
            "refusal": null
        })
    );
    Ok(())
}

#[test]
fn query_check_keeps_the_prefix_and_marks_the_whole_stage_after_count() -> TestResult {
    assert_eq!(
        answer(
            "query.check",
            &json!({"query": "patients | count | limit(1)"})
        )?,
        json!({
            "stages": [
                {"word": "patients", "span": {"start": 0, "end": 8}},
                {"word": "count", "span": {"start": 11, "end": 16}}
            ],
            "refusal": {
                "code": "query_parse",
                "message": "nothing follows count at bytes 19..27",
                "span": {"start": 19, "end": 27}
            }
        })
    );
    Ok(())
}

#[test]
fn query_check_counts_utf8_bytes_in_the_text_as_sent() -> TestResult {
    let query = " \u{2003}collection(\"café\") | 'é | .' | limit(2) \n";
    let answer = answer("query.check", &json!({"query": query}))?;
    let stages = answer["stages"].as_array().ok_or("stages missing")?;
    assert_eq!(stages.len(), 3);
    for (stage, written) in stages
        .iter()
        .zip(["collection(\"café\")", "'é | .'", "limit(2)"])
    {
        let start = query.find(written).ok_or("stage missing")?;
        assert_eq!(
            stage["span"],
            json!({"start": start, "end": start + written.len()})
        );
    }
    assert_eq!(stages[0]["word"], "collection");
    assert_eq!(stages[1]["word"], "jq");
    assert!(answer["refusal"].is_null());
    Ok(())
}

#[test]
fn query_check_keeps_completed_stages_when_arguments_are_invalid() -> TestResult {
    let checked = answer(
        "query.check",
        &json!({"query": "patients | sort(.a) | limit(0)"}),
    )?;
    assert_eq!(
        checked["stages"].as_array().ok_or("stages missing")?.len(),
        2
    );
    assert_eq!(checked["refusal"]["code"], "query_parse");
    assert_eq!(checked["refusal"]["span"], json!({"start": 28, "end": 29}));
    let invalid_source = answer("query.check", &json!({"query": "."}))?;
    assert_eq!(invalid_source["stages"], json!([]));
    assert_eq!(invalid_source["refusal"]["code"], "query_parse");
    Ok(())
}

#[test]
fn query_words_gives_every_served_reserved_and_jq_word() -> TestResult {
    let value = answer("query.words", &json!({}))?;
    let listed = value.as_array().ok_or("words missing")?;
    let names: BTreeSet<_> = listed
        .iter()
        .map(|word| word["name"].as_str().ok_or("name missing"))
        .collect::<Result<_, _>>()?;
    assert_eq!(names.len(), listed.len(), "duplicate words");
    for name in [
        "sort",
        "count",
        "limit",
        "similar",
        "collection",
        "record",
        "cosine",
        "l2",
        "dot",
        "rg",
        "match",
        "window",
        "rerank",
        "{",
        "group_by",
        "percentile",
    ] {
        let word = listed
            .iter()
            .find(|word| word["name"] == name)
            .ok_or("served word missing")?;
        assert_eq!(word["status"], "served", "{name}");
    }
    for (name, role) in [
        ("rerank", "order"),
        ("{", "shape"),
        ("group_by", "group"),
        ("percentile", "aggregate"),
    ] {
        let word = listed
            .iter()
            .find(|word| word["name"] == name)
            .ok_or("served word missing")?;
        assert_eq!(word["role"], role, "{name}");
    }
    {
        let name = "sort_by";
        let word = listed
            .iter()
            .find(|word| word["name"] == name)
            .ok_or("reserved word missing")?;
        assert_eq!(word["status"], "not_yet", "{name}");
        assert_eq!(word["role"], "stage-not-yet", "{name}");
    }
    for name in [
        "as", "def", "if", "then", "elif", "else", "end", "and", "or", "reduce", "foreach", "in",
        "try", "catch", "label", "break", "module", "import", "include", "true", "false", "null",
    ] {
        let word = listed
            .iter()
            .find(|word| word["name"] == name)
            .ok_or("jq keyword missing")?;
        assert_eq!(word["role"], "jq-keyword", "{name}");
        assert_eq!(word["status"], "served", "{name}");
    }
    for word in listed {
        assert!(word["arguments"].is_array());
        assert!(
            !word["meaning"]
                .as_str()
                .ok_or("meaning missing")?
                .is_empty()
        );
        assert!(word["role"].is_string());
    }
    assert_eq!(value, serde_json::to_value(words::WORDS)?);
    Ok(())
}

#[test]
fn query_words_fixture_is_exactly_the_pretty_table() -> TestResult {
    let expected = serde_json::to_string_pretty(words::WORDS)?;
    assert_eq!(
        include_str!("words.json"),
        expected,
        "expected words.json:\n{expected}"
    );
    Ok(())
}

#[test]
fn query_language_verbs_answer_without_any_store_configured() -> TestResult {
    // The production handlers take only their requests, so no store can be opened here.
    let checked = answer("query.check", &json!({"query": "patients | count"}))?;
    assert!(checked["refusal"].is_null());
    assert!(
        !answer("query.words", &json!({}))?
            .as_array()
            .ok_or("words missing")?
            .is_empty()
    );
    for method in ["query.check", "query.words"] {
        let params = if method == "query.check" {
            json!({"query": "patients"})
        } else {
            json!({})
        };
        assert!(authorize::subjects(&operation(method, &params)?)?.is_empty());
        assert!(route_matrix::ENVELOPES.contains(&method));
    }
    Ok(())
}

#[test]
fn query_language_requests_refuse_extra_or_missing_fields() {
    for (method, params) in [
        ("query.check", json!({})),
        ("query.check", json!({"query": 7})),
        ("query.check", json!({"query": "patients", "store": "ward"})),
        ("query.words", json!({"store": "ward"})),
    ] {
        assert!(operation(method, &params).is_err(), "{method}: {params}");
    }
}

#[test]
fn query_reserved_refusals_are_taken_from_the_word_table() -> TestResult {
    for word in words::WORDS
        .iter()
        .filter(|word| word.status == words::Status::NotYet)
    {
        let query = format!("patients | {}", word.name);
        let refused = parser::pipeline(&query)
            .err()
            .ok_or("reserved stage accepted")?;
        assert_eq!(refused.code(), "query_stage_not_available", "{}", word.name);
        assert_eq!(
            refused.to_string(),
            format!(
                "{} is not available yet at bytes 11..{}",
                word.refusal.ok_or("refusal missing")?,
                11 + word.name.len()
            )
        );
        let checked = answer("query.check", &json!({"query": query}))?;
        assert_eq!(
            checked["stages"].as_array().ok_or("stages missing")?.len(),
            1
        );
        assert_eq!(checked["refusal"]["message"], refused.to_string());
    }
    Ok(())
}

#[test]
fn query_check_and_pipeline_make_the_same_acceptance_decision() -> TestResult {
    for query in [
        "patients",
        "patients | sort(.a) | limit(2)",
        "patients | count | limit(1)",
        "patients | '.a' | sort(.b)",
        "patients | similar(.v; record(\"p\"); cosine)",
        "patients | rerank",
        "patients | 'unclosed",
        "patients | '.a' | '.b'",
        "",
    ] {
        let checked = answer("query.check", &json!({"query": query}))?;
        match parser::pipeline(query) {
            Ok(_) => assert!(checked["refusal"].is_null(), "{query}"),
            Err(refused) => {
                assert_eq!(checked["refusal"]["code"], refused.code(), "{query}");
                assert_eq!(
                    checked["refusal"]["message"],
                    refused.to_string(),
                    "{query}"
                );
                let span = refused.span();
                assert_eq!(
                    checked["refusal"]["span"],
                    json!({"start": span.start, "end": span.end}),
                    "{query}"
                );
            }
        }
    }
    Ok(())
}

#[test]
fn query_words_serves_grep_as_search_with_the_rg_arguments() -> TestResult {
    let answer = answer("query.words", &json!({}))?;
    let words = answer.as_array().ok_or("words missing")?;
    let position = words
        .iter()
        .position(|word| word["name"] == "rg")
        .ok_or("rg missing")?;
    let rg = words.get(position).ok_or("rg missing")?;
    let grep = words.get(position + 1).ok_or("grep missing after rg")?;
    assert_eq!(grep["name"], "grep");
    assert_eq!(grep["role"], "search");
    assert_eq!(grep["status"], "served");
    assert_eq!(grep["arguments"], rg["arguments"]);
    assert_eq!(
        grep["meaning"],
        "Alias of rg: keep whole records in order when their compact JSON or named field matches the pattern."
    );
    Ok(())
}