AkurAI-Notes / crates/notes-mcp/tests/mcp_test.rs

Latest change 705e1ce035ea4b9fb4c413dd38f550650185b571 - release: v0.0.5 by Claude Code

use akurai_storage::BTree;
use notes_api::embed::EmbedStore;
use std::io::{Read, Write};
use std::net::TcpListener;
use std::sync::{Arc, Mutex};
use std::thread;

fn make_registry() -> notes_mcp::tools::ToolRegistry {
    make_registry_with_embed(None, None, None)
}

fn make_registry_with_embed(
    embed_url: Option<String>,
    embed_model: Option<String>,
    embed_api_key: Option<String>,
) -> notes_mcp::tools::ToolRegistry {
    let dir = std::env::temp_dir().join(format!("notes-mcp-test-{}", std::process::id()));
    let _ = std::fs::create_dir_all(&dir);
    let unique = format!(
        "test-{}.db",
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_nanos()
    );
    let db_path = dir.join(unique);
    let db = BTree::open(&db_path).expect("create test db");
    let owner_sub = "testuser@example.com".to_string();
    notes_mcp::tools::ToolRegistry::new(
        Arc::new(Mutex::new(db)),
        embed_url,
        embed_model,
        embed_api_key,
        owner_sub,
    )
}

fn spawn_embed_server(expected_auth: &str) -> (String, thread::JoinHandle<String>) {
    let listener = TcpListener::bind("127.0.0.1:0").expect("bind test embed server");
    let addr = format!("http://{}", listener.local_addr().unwrap());
    let expected_auth = expected_auth.to_string();
    let handle = thread::spawn(move || {
        let (mut stream, _) = listener.accept().expect("accept test request");
        let mut request = Vec::new();
        let mut buf = [0_u8; 1024];
        loop {
            let read = stream.read(&mut buf).expect("read request");
            if read == 0 {
                break;
            }
            request.extend_from_slice(&buf[..read]);
            let text = String::from_utf8_lossy(&request);
            let Some(head_end) = text.find("\r\n\r\n") else {
                continue;
            };
            let content_length = text
                .lines()
                .find_map(|line| {
                    let (name, value) = line.split_once(':')?;
                    if name.eq_ignore_ascii_case("Content-Length") {
                        value.trim().parse::<usize>().ok()
                    } else {
                        None
                    }
                })
                .unwrap_or(0);
            if request.len() >= head_end + 4 + content_length {
                break;
            }
        }

        let request_text = String::from_utf8_lossy(&request).to_string();
        let expected = format!("\r\nAuthorization: Bearer {expected_auth}\r\n");
        assert!(
            request_text.contains(&expected),
            "missing bearer header in request: {request_text}"
        );

        let body = r#"{"data":[{"embedding":[0.9,0.1,0.0]}]}"#;
        let response = format!(
            "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
            body.len(),
            body
        );
        stream
            .write_all(response.as_bytes())
            .expect("write response");
        request_text
    });
    (addr, handle)
}

fn store_embedding(registry: &notes_mcp::tools::ToolRegistry, note_id: i64, vector: &[f32]) {
    let mut db = registry.db.lock().unwrap();
    let mut store = EmbedStore { db: &mut db };
    store.store(note_id, vector).unwrap();
}

fn create_note_via_tool(
    registry: &notes_mcp::tools::ToolRegistry,
    title: &str,
    body: &str,
    owner: &str,
) {
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs() as i64;

    let next_id = {
        let mut db = registry.db.lock().unwrap();
        let next = db
            .get(b"_meta:next_note_id")
            .ok()
            .flatten()
            .and_then(|b| String::from_utf8_lossy(&b).parse::<i64>().ok())
            .unwrap_or(1);
        db.insert(b"_meta:next_note_id", format!("{}", next + 1).as_bytes())
            .unwrap();
        db.commit().unwrap();
        next
    };

    let note = akurai_json::Value::Object(vec![
        ("id".into(), akurai_json::Value::Int(next_id)),
        ("title".into(), akurai_json::Value::Str(title.into())),
        ("body".into(), akurai_json::Value::Str(body.into())),
        ("ownerSub".into(), akurai_json::Value::Str(owner.into())),
        ("createdAt".into(), akurai_json::Value::Int(now)),
        ("updatedAt".into(), akurai_json::Value::Int(now)),
        ("archived".into(), akurai_json::Value::Bool(false)),
    ]);

    let mut key = b"notes:".to_vec();
    key.extend_from_slice(&(next_id as u64).to_be_bytes());

    let mut db = registry.db.lock().unwrap();
    db.insert(&key, note.to_json().as_bytes()).unwrap();
    db.commit().unwrap();
}

#[test]
fn initialize_returns_server_info() {
    let registry = make_registry();
    let req = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#;
    let resp = registry.handle_request(req);
    assert!(resp.contains("\"jsonrpc\":\"2.0\""));
    assert!(resp.contains("\"protocolVersion\":\"2024-11-05\""));
    assert!(resp.contains("\"akurai-notes-mcp\""));
}

#[test]
fn tools_list_returns_all_notes_tools() {
    let registry = make_registry();
    let req = r#"{"jsonrpc":"2.0","id":2,"method":"tools/list"}"#;
    let resp = registry.handle_request(req);
    assert!(resp.contains("\"search_notes\""));
    assert!(resp.contains("\"semantic_search\""));
    assert!(resp.contains("\"get_note\""));
    assert!(resp.contains("\"list_notebooks\""));
    assert!(resp.contains("\"create_note\""));
    assert!(resp.contains("\"update_note\""));
    assert!(resp.contains("\"delete_note\""));
    assert!(resp.contains("\"add_tag\""));
    assert!(resp.contains("\"remove_tag\""));
}

#[test]
fn bad_json_returns_parse_error() {
    let registry = make_registry();
    let resp = registry.handle_request("not json");
    assert!(resp.contains("\"code\":-32700"));
}

#[test]
fn search_notes_returns_owned_only() {
    let registry = make_registry();
    create_note_via_tool(
        &registry,
        "Meeting notes",
        "Discussed Q3 roadmap",
        "testuser@example.com",
    );
    create_note_via_tool(&registry, "Other note", "Not relevant", "other@example.com");

    let req = r#"{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"search_notes","arguments":{"query":"meeting"}}}"#;
    let resp = registry.handle_request(req);
    assert!(resp.contains("Meeting notes"));
}

#[test]
fn search_notes_tokenizes_and_ranks_title_match() {
    let registry = make_registry();
    create_note_via_tool(
        &registry,
        "AkurAI Platform - Skills & CLI Tools",
        "Canonical tool inventory",
        "testuser@example.com",
    );
    create_note_via_tool(
        &registry,
        "AkurAI EC2 - Documentation Index",
        "Cross-cutting docs link to AkurAI Platform Skills CLI Tools",
        "testuser@example.com",
    );

    let req = r#"{"jsonrpc":"2.0","id":30,"method":"tools/call","params":{"name":"search_notes","arguments":{"query":"AkurAI Platform Skills CLI Tools","k":2}}}"#;
    let resp = registry.handle_request(req);
    let title_pos = resp
        .find("AkurAI Platform - Skills & CLI Tools")
        .expect("title match should be returned");
    let index_pos = resp
        .find("AkurAI EC2 - Documentation Index")
        .expect("cross-reference match should be returned");
    assert!(
        title_pos < index_pos,
        "title match should rank before body-only cross-reference: {resp}"
    );
}

#[test]
fn semantic_search_without_embed_uses_lexical_query() {
    let registry = make_registry();
    create_note_via_tool(
        &registry,
        "Deploy runbook",
        "akurai-ec2 ports drift before release",
        "testuser@example.com",
    );
    create_note_via_tool(
        &registry,
        "Grocery list",
        "milk and eggs",
        "testuser@example.com",
    );

    let req = r#"{"jsonrpc":"2.0","id":31,"method":"tools/call","params":{"name":"semantic_search","arguments":{"query":"akurai ec2 ports","k":5}}}"#;
    let resp = registry.handle_request(req);
    assert!(resp.contains("Deploy runbook"), "{resp}");
    assert!(!resp.contains("Grocery list"), "{resp}");
}

#[test]
fn semantic_search_uses_embed_bearer_token() {
    let (embed_url, handle) = spawn_embed_server("router-secret");
    let registry = make_registry_with_embed(
        Some(embed_url),
        Some("test-embed-model".into()),
        Some("router-secret".into()),
    );
    create_note_via_tool(
        &registry,
        "Deploy runbook",
        "akurai-ec2 ports drift before release",
        "testuser@example.com",
    );
    store_embedding(&registry, 1, &[0.9, 0.1, 0.0]);

    let req = r#"{"jsonrpc":"2.0","id":32,"method":"tools/call","params":{"name":"semantic_search","arguments":{"query":"protected query","k":5}}}"#;
    let resp = registry.handle_request(req);

    assert!(resp.contains("Deploy runbook"), "{resp}");
    let request = handle.join().unwrap();
    assert!(request.contains(r#""model":"test-embed-model""#));
}

#[test]
fn get_note_returns_full_note() {
    let registry = make_registry();
    create_note_via_tool(
        &registry,
        "Test Note",
        "Hello world",
        "testuser@example.com",
    );

    let req = r#"{"jsonrpc":"2.0","id":4,"method":"tools/call","params":{"name":"get_note","arguments":{"id":1}}}"#;
    let resp = registry.handle_request(req);
    assert!(resp.contains("Test Note"));
    assert!(resp.contains("Hello world"));
}

#[test]
fn create_note_inserts_and_returns_id() {
    let registry = make_registry();
    let req = r#"{"jsonrpc":"2.0","id":5,"method":"tools/call","params":{"name":"create_note","arguments":{"title":"New Note","body":"Some body"}}}"#;
    let resp = registry.handle_request(req);
    assert!(resp.contains("\"result\""));

    let get_req = r#"{"jsonrpc":"2.0","id":6,"method":"tools/call","params":{"name":"get_note","arguments":{"id":1}}}"#;
    let get_resp = registry.handle_request(get_req);
    assert!(get_resp.contains("New Note"));
}

#[test]
fn unknown_tool_returns_error() {
    let registry = make_registry();
    let req = r#"{"jsonrpc":"2.0","id":7,"method":"tools/call","params":{"name":"nonexistent","arguments":{}}}"#;
    let resp = registry.handle_request(req);
    assert!(resp.contains("\"code\":-32603"));
}

#[test]
fn unknown_method_returns_error() {
    let registry = make_registry();
    let req = r#"{"jsonrpc":"2.0","id":8,"method":"bogus"}"#;
    let resp = registry.handle_request(req);
    assert!(resp.contains("\"code\":-32601"));
}

// Regression: JSON-RPC notifications (no `id`) must produce NO response body so the
// HTTP layer can answer 202 Accepted. A body here breaks strict clients (codex/rmcp):
// "Transport channel closed, when send initialized notification".
#[test]
fn notification_initialized_has_no_response_body() {
    let registry = make_registry();
    let resp = registry.handle_request(r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#);
    assert!(
        resp.is_empty(),
        "notification must yield empty body, got: {resp}"
    );
}

#[test]
fn id_less_message_has_no_response_body() {
    let registry = make_registry();
    let resp = registry.handle_request(r#"{"jsonrpc":"2.0","method":"tools/list"}"#);
    assert!(
        resp.is_empty(),
        "id-less message is a notification, got: {resp}"
    );
}