kibble 0.1.0

chew through any source into clean datasets — a fast ingestion, RAG & fine-tuning toolkit
Documentation
use std::path::Path;
use serde_json::{json, Value};

pub fn tool_schemas() -> Value {
    let path_arg = json!({"path":{"type":"string","description":"working dir (default '.')"}});
    json!([
        {"name":"build","description":"Build the training dataset from sources; returns BuildStats.","inputSchema":{"type":"object","properties":path_arg}},
        {"name":"eval","description":"Grade dataset quality; returns the QualityReport (check `overall`).","inputSchema":{"type":"object","properties":{"path":{"type":"string"}}}},
        {"name":"bench","description":"Benchmark a served model; returns the BenchReport (check `overall`).","inputSchema":{"type":"object","properties":{"path":{"type":"string"}}}},
        {"name":"soul_build","description":"Compile soul.toml into serving artifacts; returns the manifest.","inputSchema":{"type":"object","properties":{"soul":{"type":"string"},"out":{"type":"string"}}}},
        {"name":"caps_list","description":"List installed capability skills.","inputSchema":{"type":"object","properties":{"path":{"type":"string"},"project":{"type":"boolean"}}}},
        {"name":"caps_install","description":"Install skills from a path or URL.","inputSchema":{"type":"object","properties":{"src":{"type":"string"},"force":{"type":"boolean"},"project":{"type":"boolean"}},"required":["src"]}}
    ])
}

async fn call_tool(repo_root: &Path, name: &str, args: &Value) -> Result<String, String> {
    let s = |k: &str| args.get(k).and_then(|v| v.as_str());
    let b = |k: &str| args.get(k).and_then(|v| v.as_bool()).unwrap_or(false);
    let dir = |k: &str| s(k).map(|p| repo_root.join(p)).unwrap_or_else(|| repo_root.to_path_buf());
    match name {
        "build" => {
            let stats = crate::build::run_build(&dir("path")).await.map_err(|e| e.to_string())?;
            serde_json::to_string_pretty(&stats).map_err(|e| e.to_string())
        }
        "eval" => {
            let rep = crate::eval::run_eval(&dir("path"), b("strict")).await.map_err(|e| e.to_string())?;
            serde_json::to_string_pretty(&rep).map_err(|e| e.to_string())
        }
        "bench" => {
            let rep = crate::bench::run_bench(&dir("path"), b("strict"), false).await.map_err(|e| e.to_string())?;
            serde_json::to_string_pretty(&rep).map_err(|e| e.to_string())
        }
        "soul_build" => {
            let soul = repo_root.join(s("soul").unwrap_or("soul.toml"));
            let out = repo_root.join(s("out").unwrap_or("soul/dist"));
            let m = crate::soul::run_soul_build(&soul, &out).map_err(|e| e.to_string())?;
            serde_json::to_string_pretty(&m).map_err(|e| e.to_string())
        }
        "caps_list" => {
            let cfg = crate::config::load_config(&repo_root.join(crate::config::CONFIG_FILE));
            let root = crate::caps::caps_root(&cfg.caps, if b("project") { Some(repo_root) } else { None });
            serde_json::to_string_pretty(&crate::caps::list(&root)).map_err(|e| e.to_string())
        }
        "caps_install" => {
            let cfg = crate::config::load_config(&repo_root.join(crate::config::CONFIG_FILE));
            let src = s("src").ok_or_else(|| "missing 'src'".to_string())?;
            let names = crate::caps::install_source(repo_root, &cfg, src, b("force"), b("project"), false)
                .await.map_err(|e| e.to_string())?;
            serde_json::to_string_pretty(&names).map_err(|e| e.to_string())
        }
        other => Err(format!("unknown tool '{other}'")),
    }
}

fn err_response(id: Value, code: i64, message: &str) -> Value {
    json!({"jsonrpc":"2.0","id":id,"error":{"code":code,"message":message}})
}

pub async fn handle_message(repo_root: &Path, msg: &Value) -> Option<Value> {
    let id = msg.get("id").cloned();
    let id = id?;                          // no id → notification → no response
    let method = msg.get("method").and_then(|m| m.as_str()).unwrap_or("");
    match method {
        "initialize" => {
            let pv = msg.get("params").and_then(|p| p.get("protocolVersion")).and_then(|v| v.as_str()).unwrap_or("2024-11-05");
            Some(json!({"jsonrpc":"2.0","id":id,"result":{
                "protocolVersion": pv,
                "capabilities": {"tools": {}},
                "serverInfo": {"name":"kibble","version": env!("CARGO_PKG_VERSION")}
            }}))
        }
        "ping" => Some(json!({"jsonrpc":"2.0","id":id,"result":{}})),
        "tools/list" => Some(json!({"jsonrpc":"2.0","id":id,"result":{"tools": tool_schemas()}})),
        "tools/call" => {
            let params = msg.get("params").cloned().unwrap_or_else(|| json!({}));
            let name = params.get("name").and_then(|n| n.as_str()).unwrap_or("");
            let args = params.get("arguments").cloned().unwrap_or_else(|| json!({}));
            // required-arg validation → -32602
            if name == "caps_install" && args.get("src").and_then(|s| s.as_str()).filter(|s| !s.is_empty()).is_none() {
                return Some(err_response(id, -32602, "caps_install requires 'src'"));
            }
            let result = match call_tool(repo_root, name, &args).await {
                Ok(text) => json!({"content":[{"type":"text","text":text}], "isError": false}),
                Err(e) => json!({"content":[{"type":"text","text":e}], "isError": true}),
            };
            Some(json!({"jsonrpc":"2.0","id":id,"result":result}))
        }
        _ => Some(err_response(id, -32601, "method not found")),
    }
}

pub async fn run_mcp(repo_root: &Path) -> std::io::Result<()> {
    use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
    let mut lines = BufReader::new(tokio::io::stdin()).lines();
    let mut out = tokio::io::stdout();
    while let Some(line) = lines.next_line().await? {
        let trimmed = line.trim();
        if trimmed.is_empty() { continue; }
        let response = match serde_json::from_str::<Value>(trimmed) {
            Ok(msg) => handle_message(repo_root, &msg).await,
            Err(_) => Some(json!({"jsonrpc":"2.0","id":null,"error":{"code":-32700,"message":"parse error"}})),
        };
        if let Some(resp) = response {
            out.write_all(resp.to_string().as_bytes()).await?;
            out.write_all(b"\n").await?;
            out.flush().await?;
        }
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn initialize_and_toolslist_and_ping() {
        let root = Path::new(".");
        let init = handle_message(root, &json!({"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18"}})).await.unwrap();
        assert_eq!(init["result"]["serverInfo"]["name"], "kibble");
        assert_eq!(init["result"]["protocolVersion"], "2025-06-18");   // echoed
        assert!(init["result"]["capabilities"]["tools"].is_object());

        let list = handle_message(root, &json!({"jsonrpc":"2.0","id":2,"method":"tools/list"})).await.unwrap();
        let names: Vec<&str> = list["result"]["tools"].as_array().unwrap().iter().map(|t| t["name"].as_str().unwrap()).collect();
        for want in ["build","eval","bench","soul_build","caps_list","caps_install"] {
            assert!(names.contains(&want), "missing tool {want}");
        }
        assert!(list["result"]["tools"][0]["inputSchema"].is_object());

        let ping = handle_message(root, &json!({"jsonrpc":"2.0","id":3,"method":"ping"})).await.unwrap();
        assert!(ping["result"].is_object());
    }

    #[tokio::test]
    async fn notifications_and_unknown_and_badargs() {
        let root = Path::new(".");
        // notification (no id) → no response
        assert!(handle_message(root, &json!({"jsonrpc":"2.0","method":"notifications/initialized"})).await.is_none());
        // unknown method → -32601
        let unk = handle_message(root, &json!({"jsonrpc":"2.0","id":9,"method":"nope"})).await.unwrap();
        assert_eq!(unk["error"]["code"], -32601);
        // caps_install missing src → -32602
        let bad = handle_message(root, &json!({"jsonrpc":"2.0","id":10,"method":"tools/call","params":{"name":"caps_install","arguments":{}}})).await.unwrap();
        assert_eq!(bad["error"]["code"], -32602);
    }

    #[tokio::test]
    async fn call_eval_and_caps_list() {
        // temp repo with a minimal built dataset
        let root = std::env::temp_dir().join(format!("kibble_mcp_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&root);
        let ds = root.join("data/datasets/unsloth");
        std::fs::create_dir_all(&ds).unwrap();
        let row = |u:&str,a:&str| format!("{{\"messages\":[{{\"role\":\"user\",\"content\":{}}},{{\"role\":\"assistant\",\"content\":{}}}]}}", serde_json::to_string(u).unwrap(), serde_json::to_string(a).unwrap());
        let train: String = (0..60).map(|i| row(&format!("q{i}"), &format!("distinct answer {i} alpha beta gamma delta epsilon zeta eta theta iota kappa number {i}"))).collect::<Vec<_>>().join("\n");
        std::fs::write(ds.join("train.jsonl"), train).unwrap();
        std::fs::write(ds.join("valid.jsonl"), row("held","a fresh held out answer alpha beta gamma delta epsilon")).unwrap();
        std::fs::write(ds.join("test.jsonl"), "").unwrap();
        std::fs::write(root.join(crate::config::CONFIG_FILE), "[paths]\ndataset_dir=\"data/datasets/unsloth\"\n").unwrap();

        // eval tool → content has a quality report, not an error
        let ev = handle_message(&root, &json!({"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"eval","arguments":{}}})).await.unwrap();
        assert_eq!(ev["result"]["isError"], false);
        let text = ev["result"]["content"][0]["text"].as_str().unwrap();
        assert!(text.contains("quality_score"));

        // caps_list tool (project scope, empty registry) → not an error, valid JSON array
        let cl = handle_message(&root, &json!({"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"caps_list","arguments":{"project":true}}})).await.unwrap();
        assert_eq!(cl["result"]["isError"], false);

        // unknown tool → isError true
        let bad = handle_message(&root, &json!({"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"nope","arguments":{}}})).await.unwrap();
        assert_eq!(bad["result"]["isError"], true);
    }
}