Menu
Public repository / master
AkurAI-Notes / crates/notes-mcp/src/tools.rs
Latest change 93c3285990a1755480fbde74cd70a80bc4ee738f - fix: format lean cleanup (AKURAI-NOTES-10 AKURAI-NOTES-11 AKURAI-NOTES-12) by Ólafur Búi Ólafsson
use akurai_json::Value;
use akurai_storage::BTree;
use notes_api::embed::{embed_query_with_key, maybe_reembed, EmbedStore};
use notes_api::handlers::crud::allocate_next_id;
use notes_core::search as note_search;
use std::collections::HashSet;
use std::sync::mpsc::{channel, Sender};
use std::sync::{Arc, Mutex};
use std::thread;
/// A queued re-embedding job: recompute and store the embedding vector for a
/// note out of band, so request handlers never hold the database lock across
/// the network-bound embedding call.
struct EmbedJob {
note_id: i64,
body: String,
}
pub struct ToolRegistry {
pub db: Arc<Mutex<BTree>>,
pub embed_url: Option<String>,
pub embed_model: Option<String>,
pub embed_api_key: Option<String>,
pub owner_sub: String,
/// When present, note create/update enqueue re-embedding here instead of
/// running it inline. The background worker embeds without holding the db
/// lock, so concurrent requests queue rather than time out. `None` (the CLI
/// one-shot path) falls back to synchronous embedding.
embed_tx: Option<Sender<EmbedJob>>,
}
impl ToolRegistry {
/// Synchronous-embedding registry (one-shot CLI use): embedding runs inline
/// on the calling thread so it completes before the process exits.
pub fn new(
db: Arc<Mutex<BTree>>,
embed_url: Option<String>,
embed_model: Option<String>,
embed_api_key: Option<String>,
owner_sub: String,
) -> Self {
Self {
db,
embed_url,
embed_model,
embed_api_key,
owner_sub,
embed_tx: None,
}
}
/// Long-running-server registry: spawns a background worker that drains a
/// FIFO queue of re-embedding jobs. Note create/update persist under the db
/// lock, enqueue the job, and return immediately; the worker performs the
/// network embedding without the lock and only re-acquires it briefly to
/// store the vector. This keeps every other MCP request responsive while an
/// embedding is in flight instead of starving them into 30s timeouts.
pub fn new_background(
db: Arc<Mutex<BTree>>,
embed_url: Option<String>,
embed_model: Option<String>,
embed_api_key: Option<String>,
owner_sub: String,
) -> Self {
let embed_tx = spawn_embed_worker(
Arc::clone(&db),
embed_url.clone(),
embed_model.clone(),
embed_api_key.clone(),
);
Self {
db,
embed_url,
embed_model,
embed_api_key,
owner_sub,
embed_tx: Some(embed_tx),
}
}
/// Schedule re-embedding for a note. Uses the background queue when present
/// (server), otherwise embeds synchronously (CLI one-shot). Callers hold the
/// db lock; queuing never performs the network call under that lock.
fn schedule_reembed(&self, db: &mut BTree, note_id: i64, body: String) {
match &self.embed_tx {
Some(tx) => {
if let Err(err) = tx.send(EmbedJob { note_id, body }) {
// Worker gone: fall back to inline embedding so the vector
// is not silently lost.
let EmbedJob { note_id, body } = err.0;
maybe_reembed(
db,
&self.embed_url,
&self.embed_model,
&self.embed_api_key,
note_id,
&body,
);
}
}
None => maybe_reembed(
db,
&self.embed_url,
&self.embed_model,
&self.embed_api_key,
note_id,
&body,
),
}
}
fn owner_sub(&self) -> &str {
&self.owner_sub
}
pub fn list_tools(&self) -> Vec<Value> {
vec![
self.tool_def(
"search_notes",
"Search notes by query (semantic + substring hybrid)",
vec![
("query", "string", "Search query", true),
("k", "number", "Max results (default 20)", false),
],
),
self.tool_def(
"semantic_search",
"Search notes using semantic embedding only",
vec![
("query", "string", "Search query", true),
("k", "number", "Max results (default 20)", false),
],
),
self.tool_def(
"get_note",
"Get a single note by ID",
vec![("id", "number", "Note ID", true)],
),
self.tool_def(
"list_notes",
"List all notes, optionally filtered by notebookId",
vec![("notebookId", "number", "Filter by notebook ID", false)],
),
self.tool_def("list_notebooks", "List all notebooks", vec![]),
self.tool_def(
"create_note",
"Create a new note",
vec![
("title", "string", "Note title", true),
("body", "string", "Note body (markdown)", false),
("notebookId", "number", "Notebook ID (default 1)", false),
("archived", "boolean", "Start archived", false),
],
),
self.tool_def(
"update_note",
"Update an existing note",
vec![
("id", "number", "Note ID", true),
("title", "string", "New title", false),
("body", "string", "New body", false),
("notebookId", "number", "Move to notebook", false),
("archived", "boolean", "Archive or unarchive", false),
],
),
self.tool_def(
"delete_note",
"Delete a note",
vec![("id", "number", "Note ID", true)],
),
self.tool_def(
"create_notebook",
"Create a new notebook (folder)",
vec![
("name", "string", "Notebook name", true),
("description", "string", "Notebook description", false),
("color", "string", "Hex color like #5b8cff", false),
],
),
self.tool_def(
"update_notebook",
"Update an existing notebook",
vec![
("id", "number", "Notebook ID", true),
("name", "string", "New name", false),
("description", "string", "New description", false),
("color", "string", "New hex color", false),
],
),
self.tool_def(
"move_notebook",
"Move a notebook under another notebook and set its sort order",
vec![
("id", "number", "Notebook ID", true),
("parentId", "number", "Destination parent notebook ID", true),
("sortOrder", "number", "Sort order among siblings", false),
],
),
self.tool_def(
"delete_notebook",
"Delete a notebook",
vec![("id", "number", "Notebook ID", true)],
),
self.tool_def(
"get_notebook",
"Get a single notebook by ID",
vec![("id", "number", "Notebook ID", true)],
),
self.tool_def(
"create_tag",
"Create a new tag",
vec![
("name", "string", "Tag name", true),
("color", "string", "Hex color like #7af0d0", false),
],
),
self.tool_def("list_tags", "List all tags", vec![]),
self.tool_def(
"update_tag",
"Update an existing tag",
vec![
("id", "number", "Tag ID", true),
("name", "string", "New name", false),
("color", "string", "New hex color", false),
],
),
self.tool_def(
"delete_tag",
"Delete a tag",
vec![("id", "number", "Tag ID", true)],
),
self.tool_def(
"get_tag",
"Get a single tag by ID",
vec![("id", "number", "Tag ID", true)],
),
self.tool_def(
"add_tag",
"Add a tag to a note",
vec![
("noteId", "number", "Note ID", true),
("tagId", "number", "Tag ID", true),
],
),
self.tool_def(
"remove_tag",
"Remove a tag from a note",
vec![
("noteId", "number", "Note ID", true),
("tagId", "number", "Tag ID", true),
],
),
self.tool_def(
"system_status",
"Report Notes storage health and owner-scoped record counts",
vec![],
),
]
}
fn tool_def(
&self,
name: &str,
description: &str,
params: Vec<(&str, &str, &str, bool)>,
) -> Value {
let mut properties = Vec::new();
let mut required = Vec::new();
for (pname, ptype, pdesc, is_required) in params {
properties.push((
pname.to_string(),
Value::Object(vec![
("type".into(), Value::Str(ptype.into())),
("description".into(), Value::Str(pdesc.into())),
]),
));
if is_required {
required.push(Value::Str(pname.to_string()));
}
}
Value::Object(vec![
("name".into(), Value::Str(name.into())),
("description".into(), Value::Str(description.into())),
(
"inputSchema".into(),
Value::Object(vec![
("type".into(), Value::Str("object".into())),
("properties".into(), Value::Object(properties)),
("required".into(), Value::Array(required)),
]),
),
])
}
pub fn execute(&self, tool_name: &str, args: &[(String, Value)]) -> Result<Value, String> {
match tool_name {
"search_notes" => self.search_notes(args),
"semantic_search" => self.semantic_search(args),
"get_note" => self.get_note(args),
"list_notes" => self.list_notes(args),
"list_notebooks" => self.list_notebooks(),
"create_note" => self.create_note(args),
"update_note" => self.update_note(args),
"delete_note" => self.delete_note(args),
"create_notebook" => self.create_notebook(args),
"update_notebook" => self.update_notebook(args),
"move_notebook" => self.move_notebook(args),
"delete_notebook" => self.delete_notebook(args),
"get_notebook" => self.get_notebook(args),
"create_tag" => self.create_tag(args),
"list_tags" => self.list_tags(),
"update_tag" => self.update_tag(args),
"delete_tag" => self.delete_tag(args),
"get_tag" => self.get_tag(args),
"add_tag" => self.add_tag(args),
"remove_tag" => self.remove_tag(args),
"system_status" => self.system_status(),
_ => Err(format!("unknown tool: {tool_name}")),
}
}
fn system_status(&self) -> Result<Value, String> {
let notes = match self.list_notes(&[])? {
Value::Array(values) => values.len() as i64,
_ => 0,
};
let notebooks = match self.list_notebooks()? {
Value::Array(values) => values.len() as i64,
_ => 0,
};
let tags = match self.list_tags()? {
Value::Array(values) => values.len() as i64,
_ => 0,
};
Ok(Value::Object(vec![
("status".into(), Value::Str("ok".into())),
("notes".into(), Value::Int(notes)),
("notebooks".into(), Value::Int(notebooks)),
("tags".into(), Value::Int(tags)),
]))
}
pub fn handle_request(&self, body: &str) -> String {
let req = match akurai_json::parse(body) {
Ok(v) => v,
Err(e) => return self.json_rpc_error(None, -32700, &format!("Parse error: {e}")),
};
let method = req.get("method").and_then(|v| v.as_str()).unwrap_or("");
let id = req.get("id").cloned();
// JSON-RPC notifications carry no `id` and MUST NOT receive a response body.
// Return an empty string; the HTTP layer turns this into `202 Accepted` with no
// body (per the MCP Streamable-HTTP spec). Strict clients (codex/rmcp) error out
// if a notification gets a JSON-RPC response.
if id.is_none() || method.starts_with("notifications/") {
return String::new();
}
match method {
"tools/list" => {
let tools_list = self.list_tools();
self.json_rpc_result(
id,
Value::Object(vec![("tools".into(), Value::Array(tools_list))]),
)
}
"tools/call" => {
let params = match req.get("params") {
Some(Value::Object(pairs)) => pairs.clone(),
_ => vec![],
};
let tool_name = params
.iter()
.find(|(k, _)| k == "name")
.and_then(|(_, v)| v.as_str())
.unwrap_or("");
let arguments: Vec<(String, Value)> = params
.iter()
.find(|(k, _)| k == "arguments")
.and_then(|(_, v)| match v {
Value::Object(pairs) => Some(pairs.clone()),
_ => None,
})
.unwrap_or_default();
match self.execute(tool_name, &arguments) {
Ok(result) => self.json_rpc_result(
id,
Value::Object(vec![(
"content".into(),
Value::Array(vec![Value::Object(vec![
("type".into(), Value::Str("text".into())),
("text".into(), Value::Str(result.to_json())),
])]),
)]),
),
Err(e) => self.json_rpc_error(id, -32603, &format!("Tool error: {e}")),
}
}
"tools/get" => {
let params = match req.get("params") {
Some(Value::Object(pairs)) => pairs.clone(),
_ => vec![],
};
let tool_name = params
.iter()
.find(|(k, _)| k == "name")
.and_then(|(_, v)| v.as_str())
.unwrap_or("");
let tools_list = self.list_tools();
let tool = tools_list
.into_iter()
.find(|t| t.get("name").and_then(|v| v.as_str()) == Some(tool_name));
match tool {
Some(t) => self.json_rpc_result(id, Value::Object(vec![("tool".into(), t)])),
None => self.json_rpc_error(id, -32602, &format!("Unknown tool: {tool_name}")),
}
}
"initialize" => self.json_rpc_result(
id,
Value::Object(vec![
("protocolVersion".into(), Value::Str("2024-11-05".into())),
(
"capabilities".into(),
Value::Object(vec![("tools".into(), Value::Object(vec![]))]),
),
(
"serverInfo".into(),
Value::Object(vec![
("name".into(), Value::Str("akurai-notes-mcp".into())),
("version".into(), Value::Str("0.0.1".into())),
]),
),
]),
),
_ => self.json_rpc_error(id, -32601, &format!("Method not found: {method}")),
}
}
fn json_rpc_result(&self, id: Option<Value>, result: Value) -> String {
let mut resp = Vec::new();
resp.push(("jsonrpc".into(), Value::Str("2.0".into())));
resp.push(("result".into(), result));
if let Some(id_val) = id {
resp.push(("id".into(), id_val));
}
Value::Object(resp).to_json()
}
fn json_rpc_error(&self, id: Option<Value>, code: i64, message: &str) -> String {
let mut resp = Vec::new();
resp.push(("jsonrpc".into(), Value::Str("2.0".into())));
resp.push((
"error".into(),
Value::Object(vec![
("code".into(), Value::Int(code)),
("message".into(), Value::Str(message.into())),
]),
));
if let Some(id_val) = id {
resp.push(("id".into(), id_val));
}
Value::Object(resp).to_json()
}
fn get_arg<'a>(args: &'a [(String, Value)], name: &str) -> Option<&'a Value> {
args.iter().find(|(k, _)| k == name).map(|(_, v)| v)
}
fn scan_collection(&self, collection: &str) -> Result<Vec<Value>, String> {
let prefix = format!("{}:", collection);
let start = prefix.as_bytes().to_vec();
let end = upper_bound(&start);
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let mut results = Vec::new();
if let Ok(entries) = db.range(&start, &end) {
for (_, val_bytes) in entries {
let s = String::from_utf8_lossy(&val_bytes).to_string();
if let Ok(v) = akurai_json::parse(&s) {
results.push(v);
}
}
}
Ok(results)
}
fn get_record(&self, collection: &str, id: u64) -> Result<Option<Value>, String> {
let mut key = format!("{}:", collection).into_bytes();
key.extend_from_slice(&id.to_be_bytes());
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
match db.get(&key).map_err(|e| format!("storage: {e}"))? {
Some(bytes) => {
let s = String::from_utf8_lossy(&bytes).to_string();
akurai_json::parse(&s)
.map(Some)
.map_err(|e| format!("parse: {e}"))
}
None => Ok(None),
}
}
fn search_notes(&self, args: &[(String, Value)]) -> Result<Value, String> {
let query = Self::get_arg(args, "query")
.and_then(|v| v.as_str())
.unwrap_or("");
if query.is_empty() {
return Ok(Value::Array(vec![]));
}
let k = Self::get_arg(args, "k")
.and_then(|v| v.as_f64())
.map(|f| f as usize)
.unwrap_or(20);
self.lexical_note_results(query, k).map(Value::Array)
}
fn semantic_search(&self, args: &[(String, Value)]) -> Result<Value, String> {
let query = Self::get_arg(args, "query")
.and_then(|v| v.as_str())
.unwrap_or("");
if query.is_empty() {
return Ok(Value::Array(vec![]));
}
let k = Self::get_arg(args, "k")
.and_then(|v| v.as_f64())
.map(|f| f as usize)
.unwrap_or(20);
if let Some(ref url) = self.embed_url {
let model = self
.embed_model
.as_deref()
.unwrap_or("qwen3-embedding-0.6b");
if let Ok(qv) = embed_query_with_key(query, url, model, self.embed_api_key.as_deref()) {
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let mut store = EmbedStore { db: &mut db };
let owners = [self.owner_sub()];
let ranked = store.rank_notes(&qv, &owners, k);
let mut results = Vec::new();
for (id, score) in ranked {
let mut key = b"notes:".to_vec();
key.extend_from_slice(&(id as u64).to_be_bytes());
if let Ok(Some(bytes)) = store.db.get(&key) {
if let Ok(mut note) = akurai_json::parse(&String::from_utf8_lossy(&bytes)) {
if let Value::Object(ref mut fields) = note {
set_field(fields, "score", Value::Float(score as f64));
}
results.push(note);
}
}
}
if !results.is_empty() {
return Ok(Value::Array(results));
}
}
}
self.lexical_note_results(query, k).map(Value::Array)
}
fn lexical_note_results(&self, query: &str, k: usize) -> Result<Vec<Value>, String> {
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let mut scored: Vec<(Value, f32, i64, i64)> = Vec::new();
let prefix = b"notes:";
let end = upper_bound(prefix);
if let Ok(entries) = db.range(prefix, &end) {
for (_, val_bytes) in entries {
let s = String::from_utf8_lossy(&val_bytes).to_string();
if let Ok(v) = akurai_json::parse(&s) {
if !note_search::owner_matches(&v, &[self.owner_sub()]) {
continue;
}
let title = note_search::note_title(&v);
let body = note_search::note_body(&v);
if let Some(score) = note_search::lexical_score(query, title, body) {
let updated_at = note_search::note_updated_at(&v);
let id = note_search::note_id(&v);
scored.push((v, score, updated_at, id));
}
}
}
}
scored.sort_by(|a, b| {
b.1.partial_cmp(&a.1)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| b.2.cmp(&a.2))
.then_with(|| b.3.cmp(&a.3))
});
scored.truncate(k);
Ok(scored.into_iter().map(|(note, _, _, _)| note).collect())
}
fn get_note(&self, args: &[(String, Value)]) -> Result<Value, String> {
let id = Self::get_arg(args, "id")
.and_then(|v| v.as_f64())
.map(|f| f as u64)
.ok_or("missing id")?;
self.get_record("notes", id)
.map(|opt| opt.unwrap_or(Value::Null))
}
fn list_notes(&self, args: &[(String, Value)]) -> Result<Value, String> {
let filter_notebook =
Self::get_arg(args, "notebookId").and_then(|v| v.as_f64().map(|f| f as i64));
let mut notes = self.scan_collection("notes")?;
notes.retain(|n| n.get("ownerSub").and_then(|o| o.as_str()) == Some(self.owner_sub()));
if let Some(nbid) = filter_notebook {
notes.retain(|n| n.get("notebookId").and_then(|v| v.as_i64()) == Some(nbid));
}
Ok(Value::Array(notes))
}
fn list_notebooks(&self) -> Result<Value, String> {
let mut notebooks = self.scan_collection("notebooks")?;
notebooks.retain(|n| n.get("ownerSub").and_then(|o| o.as_str()) == Some(self.owner_sub()));
Ok(Value::Array(notebooks))
}
fn next_id(&self, collection: &str) -> Result<u64, String> {
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
allocate_next_id(&mut db, collection)
}
fn create_notebook(&self, args: &[(String, Value)]) -> Result<Value, String> {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs() as i64;
let id = self.next_id("notebooks")?;
let name = Self::get_arg(args, "name")
.and_then(|v| v.as_str())
.unwrap_or("Untitled Notebook");
let description = Self::get_arg(args, "description")
.and_then(|v| v.as_str())
.unwrap_or("");
let color = Self::get_arg(args, "color")
.and_then(|v| v.as_str())
.unwrap_or("#5b8cff");
let fields = vec![
("id".into(), Value::Int(id as i64)),
("name".into(), Value::Str(name.into())),
("description".into(), Value::Str(description.into())),
("ownerSub".into(), Value::Str(self.owner_sub().into())),
("color".into(), Value::Str(color.into())),
("createdAt".into(), Value::Int(now)),
];
let record = Value::Object(fields);
let mut key = Vec::from(b"notebooks:" as &[u8]);
key.extend_from_slice(&id.to_be_bytes());
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
db.insert(&key, record.to_json().as_bytes())
.map_err(|e| format!("insert: {e}"))?;
db.commit().map_err(|e| format!("commit: {e}"))?;
Ok(Value::Object(vec![
("id".into(), Value::Int(id as i64)),
("name".into(), Value::Str(name.into())),
]))
}
fn update_notebook(&self, args: &[(String, Value)]) -> Result<Value, String> {
let id = Self::get_arg(args, "id")
.and_then(|v| v.as_f64())
.map(|f| f as u64)
.ok_or("missing id")?;
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let mut key = b"notebooks:".to_vec();
key.extend_from_slice(&id.to_be_bytes());
let mut nb = match db.get(&key).map_err(|e| format!("storage: {e}"))? {
Some(bytes) => {
let s = String::from_utf8_lossy(&bytes).to_string();
akurai_json::parse(&s).map_err(|e| format!("parse: {e}"))?
}
None => return Err(format!("notebook {id} not found")),
};
if nb.get("ownerSub").and_then(|o| o.as_str()) != Some(self.owner_sub()) {
return Err("unauthorized".into());
}
if let Some(name) = Self::get_arg(args, "name").and_then(|v| v.as_str()) {
if let Value::Object(ref mut pairs) = nb {
set_field(pairs, "name", Value::Str(name.to_string()));
}
}
if let Some(desc) = Self::get_arg(args, "description").and_then(|v| v.as_str()) {
if let Value::Object(ref mut pairs) = nb {
set_field(pairs, "description", Value::Str(desc.to_string()));
}
}
if let Some(color) = Self::get_arg(args, "color").and_then(|v| v.as_str()) {
if let Value::Object(ref mut pairs) = nb {
set_field(pairs, "color", Value::Str(color.to_string()));
}
}
db.insert(&key, nb.to_json().as_bytes())
.map_err(|e| format!("insert: {e}"))?;
db.commit().map_err(|e| format!("commit: {e}"))?;
Ok(Value::Object(vec![("updated".into(), Value::Bool(true))]))
}
fn move_notebook(&self, args: &[(String, Value)]) -> Result<Value, String> {
let id = Self::get_arg(args, "id")
.and_then(|value| value.as_f64())
.map(|value| value as u64)
.ok_or("missing id")?;
let parent_id = Self::get_arg(args, "parentId")
.and_then(|value| value.as_f64())
.map(|value| value as u64)
.ok_or("missing parentId")?;
let sort_order = Self::get_arg(args, "sortOrder")
.and_then(|value| value.as_f64())
.map(|value| value as i64);
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let mut key = b"notebooks:".to_vec();
key.extend_from_slice(&id.to_be_bytes());
let mut notebook = db
.get(&key)
.map_err(|e| format!("storage: {e}"))?
.ok_or_else(|| format!("notebook {id} not found"))
.and_then(|bytes| {
akurai_json::parse(&String::from_utf8_lossy(&bytes))
.map_err(|e| format!("parse: {e}"))
})?;
if notebook.get("ownerSub").and_then(|value| value.as_str()) != Some(self.owner_sub()) {
return Err("unauthorized".into());
}
let mut cursor = Some(parent_id);
let mut visited = HashSet::new();
while let Some(current) = cursor {
if current == id || !visited.insert(current) {
return Err("notebook hierarchy cycle".into());
}
let mut parent_key = b"notebooks:".to_vec();
parent_key.extend_from_slice(¤t.to_be_bytes());
let parent = db
.get(&parent_key)
.map_err(|e| format!("storage: {e}"))?
.ok_or_else(|| format!("parent notebook {current} not found"))
.and_then(|bytes| {
akurai_json::parse(&String::from_utf8_lossy(&bytes))
.map_err(|e| format!("parse: {e}"))
})?;
if parent.get("ownerSub").and_then(|value| value.as_str()) != Some(self.owner_sub()) {
return Err("unauthorized parent notebook".into());
}
cursor = parent
.get("parentId")
.and_then(|value| value.as_f64())
.map(|value| value as u64);
}
if let Value::Object(pairs) = &mut notebook {
set_field(pairs, "parentId", Value::Int(parent_id as i64));
if let Some(order) = sort_order {
set_field(pairs, "sortOrder", Value::Int(order));
}
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs() as i64;
set_field(pairs, "updatedAt", Value::Int(now));
}
db.insert(&key, notebook.to_json().as_bytes())
.map_err(|e| format!("insert: {e}"))?;
db.commit().map_err(|e| format!("commit: {e}"))?;
Ok(Value::Object(vec![("moved".into(), Value::Bool(true))]))
}
fn delete_notebook(&self, args: &[(String, Value)]) -> Result<Value, String> {
let id = Self::get_arg(args, "id")
.and_then(|v| v.as_f64())
.map(|f| f as u64)
.ok_or("missing id")?;
if id == 1 {
return Err("cannot delete the default notebook".into());
}
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let mut key = b"notebooks:".to_vec();
key.extend_from_slice(&id.to_be_bytes());
let owned = match db.get(&key).map_err(|e| format!("storage: {e}"))? {
Some(bytes) => {
let s = String::from_utf8_lossy(&bytes).to_string();
if let Ok(v) = akurai_json::parse(&s) {
v.get("ownerSub").and_then(|o| o.as_str()) == Some(self.owner_sub())
} else {
false
}
}
None => false,
};
if !owned {
return Err(format!("notebook {id} not found or unauthorized"));
}
db.delete(&key).map_err(|e| format!("delete: {e}"))?;
db.commit().map_err(|e| format!("commit: {e}"))?;
Ok(Value::Object(vec![("deleted".into(), Value::Bool(true))]))
}
fn get_notebook(&self, args: &[(String, Value)]) -> Result<Value, String> {
let id = Self::get_arg(args, "id")
.and_then(|v| v.as_f64())
.map(|f| f as u64)
.ok_or("missing id")?;
self.get_record("notebooks", id)
.map(|opt| opt.unwrap_or(Value::Null))
}
fn create_tag(&self, args: &[(String, Value)]) -> Result<Value, String> {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs() as i64;
let id = self.next_id("tags")?;
let name = Self::get_arg(args, "name")
.and_then(|v| v.as_str())
.unwrap_or("Untitled Tag");
let color = Self::get_arg(args, "color")
.and_then(|v| v.as_str())
.unwrap_or("#7af0d0");
let fields = vec![
("id".into(), Value::Int(id as i64)),
("name".into(), Value::Str(name.into())),
("ownerSub".into(), Value::Str(self.owner_sub().into())),
("color".into(), Value::Str(color.into())),
("createdAt".into(), Value::Int(now)),
];
let record = Value::Object(fields);
let mut key = Vec::from(b"tags:" as &[u8]);
key.extend_from_slice(&id.to_be_bytes());
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
db.insert(&key, record.to_json().as_bytes())
.map_err(|e| format!("insert: {e}"))?;
db.commit().map_err(|e| format!("commit: {e}"))?;
Ok(Value::Object(vec![
("id".into(), Value::Int(id as i64)),
("name".into(), Value::Str(name.into())),
]))
}
fn list_tags(&self) -> Result<Value, String> {
let mut tags = self.scan_collection("tags")?;
tags.retain(|t| t.get("ownerSub").and_then(|o| o.as_str()) == Some(self.owner_sub()));
Ok(Value::Array(tags))
}
fn update_tag(&self, args: &[(String, Value)]) -> Result<Value, String> {
let id = Self::get_arg(args, "id")
.and_then(|v| v.as_f64())
.map(|f| f as u64)
.ok_or("missing id")?;
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let mut key = b"tags:".to_vec();
key.extend_from_slice(&id.to_be_bytes());
let mut tag = match db.get(&key).map_err(|e| format!("storage: {e}"))? {
Some(bytes) => {
let s = String::from_utf8_lossy(&bytes).to_string();
akurai_json::parse(&s).map_err(|e| format!("parse: {e}"))?
}
None => return Err(format!("tag {id} not found")),
};
if tag.get("ownerSub").and_then(|o| o.as_str()) != Some(self.owner_sub()) {
return Err("unauthorized".into());
}
if let Some(name) = Self::get_arg(args, "name").and_then(|v| v.as_str()) {
if let Value::Object(ref mut pairs) = tag {
set_field(pairs, "name", Value::Str(name.to_string()));
}
}
if let Some(color) = Self::get_arg(args, "color").and_then(|v| v.as_str()) {
if let Value::Object(ref mut pairs) = tag {
set_field(pairs, "color", Value::Str(color.to_string()));
}
}
db.insert(&key, tag.to_json().as_bytes())
.map_err(|e| format!("insert: {e}"))?;
db.commit().map_err(|e| format!("commit: {e}"))?;
Ok(Value::Object(vec![("updated".into(), Value::Bool(true))]))
}
fn delete_tag(&self, args: &[(String, Value)]) -> Result<Value, String> {
let id = Self::get_arg(args, "id")
.and_then(|v| v.as_f64())
.map(|f| f as u64)
.ok_or("missing id")?;
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let mut key = b"tags:".to_vec();
key.extend_from_slice(&id.to_be_bytes());
let owned = match db.get(&key).map_err(|e| format!("storage: {e}"))? {
Some(bytes) => {
let s = String::from_utf8_lossy(&bytes).to_string();
if let Ok(v) = akurai_json::parse(&s) {
v.get("ownerSub").and_then(|o| o.as_str()) == Some(self.owner_sub())
} else {
false
}
}
None => false,
};
if !owned {
return Err(format!("tag {id} not found or unauthorized"));
}
db.delete(&key).map_err(|e| format!("delete: {e}"))?;
db.commit().map_err(|e| format!("commit: {e}"))?;
Ok(Value::Object(vec![("deleted".into(), Value::Bool(true))]))
}
fn get_tag(&self, args: &[(String, Value)]) -> Result<Value, String> {
let id = Self::get_arg(args, "id")
.and_then(|v| v.as_f64())
.map(|f| f as u64)
.ok_or("missing id")?;
self.get_record("tags", id)
.map(|opt| opt.unwrap_or(Value::Null))
}
fn create_note(&self, args: &[(String, Value)]) -> Result<Value, String> {
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs() as i64;
let next_id = allocate_next_id(&mut db, "notes")?;
let title = Self::get_arg(args, "title")
.and_then(|v| v.as_str())
.unwrap_or("Untitled");
let body = Self::get_arg(args, "body")
.and_then(|v| v.as_str())
.unwrap_or("");
let notebook_id = Self::get_arg(args, "notebookId")
.and_then(|v| v.as_f64())
.map(|f| f as i64)
.unwrap_or(1);
let archived = Self::get_arg(args, "archived")
.and_then(|v| v.as_bool())
.unwrap_or(false);
let fields = vec![
("id".into(), Value::Int(next_id as i64)),
("title".into(), Value::Str(title.into())),
("body".into(), Value::Str(body.into())),
("format".into(), Value::Str("markdown".into())),
("notebookId".into(), Value::Int(notebook_id)),
("ownerSub".into(), Value::Str(self.owner_sub().into())),
("tagIds".into(), Value::Str("".into())),
("createdAt".into(), Value::Int(now)),
("updatedAt".into(), Value::Int(now)),
("archived".into(), Value::Bool(archived)),
];
let record = Value::Object(fields);
let mut key = Vec::from(b"notes:" as &[u8]);
key.extend_from_slice(&next_id.to_be_bytes());
db.insert(&key, record.to_json().as_bytes())
.map_err(|e| format!("insert: {e}"))?;
db.commit().map_err(|e| format!("commit: {e}"))?;
self.schedule_reembed(&mut db, next_id as i64, body.to_string());
Ok(Value::Object(vec![(
"id".into(),
Value::Int(next_id as i64),
)]))
}
fn update_note(&self, args: &[(String, Value)]) -> Result<Value, String> {
let id = Self::get_arg(args, "id")
.and_then(|v| v.as_f64())
.map(|f| f as u64)
.ok_or("missing id")?;
let new_title = Self::get_arg(args, "title").and_then(|v| v.as_str());
let new_body = Self::get_arg(args, "body").and_then(|v| v.as_str());
let new_notebook =
Self::get_arg(args, "notebookId").and_then(|v| v.as_f64().map(|f| f as i64));
let new_archived = Self::get_arg(args, "archived").and_then(|v| v.as_bool());
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let mut key = "notes:".to_string().into_bytes();
key.extend_from_slice(&id.to_be_bytes());
let mut note = match db.get(&key).map_err(|e| format!("storage: {e}"))? {
Some(bytes) => {
let s = String::from_utf8_lossy(&bytes).to_string();
akurai_json::parse(&s).map_err(|e| format!("parse: {e}"))?
}
None => return Err(format!("note {id} not found")),
};
if note.get("ownerSub").and_then(|o| o.as_str()) != Some(self.owner_sub()) {
return Err("unauthorized".into());
}
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs() as i64;
if let Value::Object(pairs) = &mut note {
if let Some(title) = new_title {
set_field(pairs, "title", Value::Str(title.to_string()));
}
if let Some(body) = new_body {
set_field(pairs, "body", Value::Str(body.to_string()));
}
if let Some(nbid) = new_notebook {
set_field(pairs, "notebookId", Value::Int(nbid));
}
if let Some(arch) = new_archived {
set_field(pairs, "archived", Value::Bool(arch));
}
set_field(pairs, "updatedAt", Value::Int(now));
}
db.insert(&key, note.to_json().as_bytes())
.map_err(|e| format!("insert: {e}"))?;
db.commit().map_err(|e| format!("commit: {e}"))?;
if new_title.is_some() || new_body.is_some() {
let body = note
.get("body")
.and_then(|value| value.as_str())
.unwrap_or("")
.to_string();
self.schedule_reembed(&mut db, id as i64, body);
}
Ok(Value::Object(vec![("updated".into(), Value::Bool(true))]))
}
fn delete_note(&self, args: &[(String, Value)]) -> Result<Value, String> {
let id = Self::get_arg(args, "id")
.and_then(|v| v.as_f64())
.map(|f| f as u64)
.ok_or("missing id")?;
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let mut key = "notes:".to_string().into_bytes();
key.extend_from_slice(&id.to_be_bytes());
let owned = match db.get(&key).map_err(|e| format!("storage: {e}"))? {
Some(bytes) => {
let s = String::from_utf8_lossy(&bytes).to_string();
if let Ok(v) = akurai_json::parse(&s) {
v.get("ownerSub").and_then(|o| o.as_str()) == Some(self.owner_sub())
} else {
false
}
}
None => false,
};
if !owned {
return Err(format!("note {id} not found or unauthorized"));
}
db.delete(&key).map_err(|e| format!("delete: {e}"))?;
db.commit().map_err(|e| format!("commit: {e}"))?;
Ok(Value::Object(vec![("deleted".into(), Value::Bool(true))]))
}
fn add_tag(&self, args: &[(String, Value)]) -> Result<Value, String> {
let note_id = Self::get_arg(args, "noteId")
.and_then(|v| v.as_f64())
.map(|f| f as u64)
.ok_or("missing noteId")?;
let tag_id = Self::get_arg(args, "tagId")
.and_then(|v| v.as_f64())
.map(|f| f as u64)
.ok_or("missing tagId")?;
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let mut key = b"notes:".to_vec();
key.extend_from_slice(¬e_id.to_be_bytes());
let mut note = match db.get(&key).map_err(|e| format!("storage: {e}"))? {
Some(bytes) => {
let s = String::from_utf8_lossy(&bytes).to_string();
akurai_json::parse(&s).map_err(|e| format!("parse: {e}"))?
}
None => return Err(format!("note {note_id} not found")),
};
if note.get("ownerSub").and_then(|o| o.as_str()) != Some(self.owner_sub()) {
return Err("unauthorized".into());
}
let tag_id_str = tag_id.to_string();
let current_ids: Vec<String> = note
.get("tagIds")
.and_then(|v| v.as_str())
.unwrap_or("")
.split(',')
.filter(|s| !s.is_empty())
.map(|s| s.to_string())
.collect();
let mut new_ids = current_ids.clone();
if !new_ids.contains(&tag_id_str) {
new_ids.push(tag_id_str);
}
if let Value::Object(ref mut pairs) = note {
set_field(pairs, "tagIds", Value::Str(new_ids.join(",")));
}
db.insert(&key, note.to_json().as_bytes())
.map_err(|e| format!("insert: {e}"))?;
db.commit().map_err(|e| format!("commit: {e}"))?;
Ok(Value::Object(vec![("ok".into(), Value::Bool(true))]))
}
fn remove_tag(&self, args: &[(String, Value)]) -> Result<Value, String> {
let note_id = Self::get_arg(args, "noteId")
.and_then(|v| v.as_f64())
.map(|f| f as u64)
.ok_or("missing noteId")?;
let tag_id = Self::get_arg(args, "tagId")
.and_then(|v| v.as_f64())
.map(|f| f as u64)
.ok_or("missing tagId")?;
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let mut key = b"notes:".to_vec();
key.extend_from_slice(¬e_id.to_be_bytes());
let mut note = match db.get(&key).map_err(|e| format!("storage: {e}"))? {
Some(bytes) => {
let s = String::from_utf8_lossy(&bytes).to_string();
akurai_json::parse(&s).map_err(|e| format!("parse: {e}"))?
}
None => return Err(format!("note {note_id} not found")),
};
if note.get("ownerSub").and_then(|o| o.as_str()) != Some(self.owner_sub()) {
return Err("unauthorized".into());
}
let tag_id_str = tag_id.to_string();
let new_ids: Vec<String> = note
.get("tagIds")
.and_then(|v| v.as_str())
.unwrap_or("")
.split(',')
.filter(|s| !s.is_empty() && *s != tag_id_str)
.map(|s| s.to_string())
.collect();
if let Value::Object(ref mut pairs) = note {
set_field(pairs, "tagIds", Value::Str(new_ids.join(",")));
}
db.insert(&key, note.to_json().as_bytes())
.map_err(|e| format!("insert: {e}"))?;
db.commit().map_err(|e| format!("commit: {e}"))?;
Ok(Value::Object(vec![("ok".into(), Value::Bool(true))]))
}
}
/// Spawn a background worker that drains re-embedding jobs off a FIFO channel.
///
/// The worker performs the network-bound embedding call WITHOUT holding the
/// database lock, then re-acquires the lock only to persist the vector. This is
/// the core of "queue instead of error": while one note is embedding, other MCP
/// requests (reads, writes, further embed enqueues) are never blocked behind the
/// embedding round-trip and cannot pile into 30s request timeouts. Jobs are
/// processed in submission order, so the last write to a note wins its vector.
///
/// The worker lives for the process lifetime (the `Sender` is held by the
/// long-running server registry). When the registry is dropped, the channel
/// closes and the loop exits.
fn spawn_embed_worker(
db: Arc<Mutex<BTree>>,
embed_url: Option<String>,
embed_model: Option<String>,
embed_api_key: Option<String>,
) -> Sender<EmbedJob> {
let (tx, rx) = channel::<EmbedJob>();
thread::spawn(move || {
let url = match embed_url {
Some(url) => url,
// No embedding configured: drain and discard so senders never block.
None => {
for _ in rx {}
return;
}
};
let model = embed_model.unwrap_or_else(|| "qwen3-embedding-0.6b".to_string());
for job in rx {
match embed_query_with_key(&job.body, &url, &model, embed_api_key.as_deref()) {
Ok(vec) => match db.lock() {
Ok(mut guard) => {
let mut store = EmbedStore { db: &mut guard };
if let Err(e) = store.store(job.note_id, &vec) {
eprintln!("store embedding for note {}: {e}", job.note_id);
}
}
Err(e) => eprintln!("embed worker db lock poisoned: {e}"),
},
Err(e) => eprintln!("re-embed failed for note {}: {e}", job.note_id),
}
}
});
tx
}
fn set_field(pairs: &mut Vec<(String, Value)>, key: &str, value: Value) {
for pair in pairs.iter_mut() {
if pair.0 == key {
pair.1 = value;
return;
}
}
pairs.push((key.to_string(), value));
}
fn upper_bound(prefix: &[u8]) -> Vec<u8> {
let mut bound = prefix.to_vec();
bound.push(0xff);
bound
}