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

Latest change 38337fdd6069cbcb0ce0327fb09aa0b474a34e76 - fix(mcp): embed notes off-request via background queue by Ólafur Búi Ólafsson

use akurai_http::Request;
use akurai_json::Value;
use akurai_storage::BTree;
use notes_api::embed::EmbedStore;
use notes_mcp::server::{McpConfig, McpServer};
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}"
    );
}

#[test]
fn shared_database_serves_mcp_without_a_second_open() {
    let db_path = std::env::temp_dir().join(format!(
        "notes-mcp-shared-{}-{}.db",
        std::process::id(),
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_nanos()
    ));
    let mut db = BTree::open(&db_path).expect("create shared test db");
    let mut key = b"notebooks:".to_vec();
    key.extend_from_slice(&1_u64.to_be_bytes());
    db.insert(
        &key,
        br##"{"id":1,"name":"Shared","ownerSub":"owner","color":"#5b8cff"}"##,
    )
    .expect("insert notebook");
    db.commit().expect("commit notebook");

    let server = McpServer::with_db(
        McpConfig {
            host: "127.0.0.1".into(),
            port: 0,
            db: db_path.display().to_string(),
            token: "test-token".into(),
            embed_url: None,
            embed_model: None,
            embed_api_key: None,
            owner_sub: "owner".into(),
        },
        Arc::new(Mutex::new(db)),
    );

    let mut request =
        Request::parse_head(b"POST /mcp HTTP/1.1\r\nAuthorization: Bearer test-token\r\n\r\n")
            .unwrap();
    request.body = br#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"list_notebooks","arguments":{}}}"#.to_vec();
    let response = server.handle_http(&request);

    assert_eq!(response.status, 200);
    assert!(String::from_utf8(response.body).unwrap().contains("Shared"));

    request.headers.clear();
    assert_eq!(server.handle_http(&request).status, 401);
}

#[test]
fn move_notebook_sets_parent_order_and_rejects_cycles() {
    let registry = make_registry();
    let parent = registry
        .execute(
            "create_notebook",
            &[("name".into(), Value::Str("Parent".into()))],
        )
        .unwrap()
        .get("id")
        .and_then(|value| value.as_i64())
        .unwrap();
    let child = registry
        .execute(
            "create_notebook",
            &[("name".into(), Value::Str("Child".into()))],
        )
        .unwrap()
        .get("id")
        .and_then(|value| value.as_i64())
        .unwrap();

    registry
        .execute(
            "move_notebook",
            &[
                ("id".into(), Value::Int(child)),
                ("parentId".into(), Value::Int(parent)),
                ("sortOrder".into(), Value::Int(10)),
            ],
        )
        .unwrap();
    let moved = registry
        .execute("get_notebook", &[("id".into(), Value::Int(child))])
        .unwrap();
    assert_eq!(
        moved.get("parentId").and_then(|value| value.as_i64()),
        Some(parent)
    );
    assert_eq!(
        moved.get("sortOrder").and_then(|value| value.as_i64()),
        Some(10)
    );

    let cycle = registry.execute(
        "move_notebook",
        &[
            ("id".into(), Value::Int(parent)),
            ("parentId".into(), Value::Int(child)),
        ],
    );
    assert_eq!(cycle.unwrap_err(), "notebook hierarchy cycle");
}

#[test]
fn create_note_persists_embedding_with_bearer_token() {
    let (embed_url, handle) = spawn_embed_server("notes-secret");
    let registry = make_registry_with_embed(
        Some(embed_url),
        Some("test-embed-model".into()),
        Some("notes-secret".into()),
    );

    let created = registry
        .execute(
            "create_note",
            &[
                ("title".into(), Value::Str("Embedded note".into())),
                ("body".into(), Value::Str("embed this body".into())),
            ],
        )
        .unwrap();

    assert_eq!(created.get("id").and_then(|value| value.as_i64()), Some(1));
    let request = handle.join().unwrap();
    assert!(request.contains(r#""model":"test-embed-model""#));
    assert!(request.contains("embed this body"));
}

#[test]
fn background_embedding_persists_out_of_band() {
    use std::time::{Duration, Instant};

    let (embed_url, handle) = spawn_embed_server("notes-secret");
    let dir = std::env::temp_dir().join(format!("notes-mcp-bg-{}", std::process::id()));
    let _ = std::fs::create_dir_all(&dir);
    let db_path = dir.join(format!(
        "bg-{}.db",
        std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_nanos()
    ));
    let db = Arc::new(Mutex::new(BTree::open(&db_path).expect("create test db")));
    let registry = notes_mcp::tools::ToolRegistry::new_background(
        Arc::clone(&db),
        Some(embed_url),
        Some("test-embed-model".into()),
        Some("notes-secret".into()),
        "testuser@example.com".into(),
    );

    // The create call returns immediately; embedding is queued, not inline.
    let created = registry
        .execute(
            "create_note",
            &[
                ("title".into(), Value::Str("Async note".into())),
                ("body".into(), Value::Str("queue this body".into())),
            ],
        )
        .unwrap();
    let id = created
        .get("id")
        .and_then(|value| value.as_i64())
        .expect("note id");

    // The embed request is made out of band by the worker thread.
    let request = handle.join().unwrap();
    assert!(request.contains("queue this body"));

    // The worker stores the vector without the request path blocking on it.
    let deadline = Instant::now() + Duration::from_secs(5);
    let vector = loop {
        {
            let mut guard = db.lock().unwrap_or_else(|poison| poison.into_inner());
            let mut store = EmbedStore { db: &mut guard };
            if let Some(vector) = store.load(id) {
                break vector;
            }
        }
        assert!(
            Instant::now() < deadline,
            "embedding was not stored in time"
        );
        thread::sleep(Duration::from_millis(25));
    };
    assert_eq!(vector, vec![0.9, 0.1, 0.0]);
}