Menu
Public repository / 9f0fd7a16e2dff7518389ff0ddbd603ebc6264ed
AkurAI-Notes / crates/notes-mcp/src/tools.rs
Latest change 21fd83f4a6eb769a0ee34aa30f721b5cdc6881af - Fix IDP OAuth: split public issuer vs internal loopback URL for token exchange by Claude Code
use akurai_json::Value;
use akurai_storage::BTree;
use std::sync::{Arc, Mutex};
pub struct ToolRegistry {
pub db: Arc<Mutex<BTree>>,
pub embed_url: Option<String>,
pub embed_model: Option<String>,
pub owner_sub: String,
}
impl ToolRegistry {
pub fn new(
db: Arc<Mutex<BTree>>,
embed_url: Option<String>,
embed_model: Option<String>,
owner_sub: String,
) -> Self {
Self {
db,
embed_url,
embed_model,
owner_sub,
}
}
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_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", 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),
],
),
self.tool_def(
"delete_note",
"Delete a note",
vec![("id", "number", "Note 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),
],
),
]
}
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_notebooks" => self.list_notebooks(),
"create_note" => self.create_note(args),
"update_note" => self.update_note(args),
"delete_note" => self.delete_note(args),
"add_tag" => self.add_tag(args),
"remove_tag" => self.remove_tag(args),
_ => Err(format!("unknown tool: {tool_name}")),
}
}
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();
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())),
]),
),
]),
),
"notifications/initialized" => self.json_rpc_result(id, Value::Null),
_ => 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("")
.to_lowercase();
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);
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let mut results = Vec::new();
let prefix = b"notes:";
let end = upper_bound(prefix);
if let Ok(entries) = db.range(prefix, &end) {
for (_, val_bytes) in entries {
if results.len() >= k {
break;
}
let s = String::from_utf8_lossy(&val_bytes).to_string();
if let Ok(v) = akurai_json::parse(&s) {
if v.get("ownerSub").and_then(|o| o.as_str()) != Some(self.owner_sub()) {
continue;
}
let title = v.get("title").and_then(|t| t.as_str()).unwrap_or("");
let body = v.get("body").and_then(|b| b.as_str()).unwrap_or("");
if title.to_lowercase().contains(&query) || body.to_lowercase().contains(&query)
{
results.push(v);
}
}
}
}
Ok(Value::Array(results))
}
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("")
.to_lowercase();
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 self.embed_url.is_some() {
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let mut results = Vec::new();
let prefix = b"notes:";
let end = upper_bound(prefix);
if let Ok(entries) = db.range(prefix, &end) {
for (_, val_bytes) in entries {
if results.len() >= k {
break;
}
let s = String::from_utf8_lossy(&val_bytes).to_string();
if let Ok(v) = akurai_json::parse(&s) {
if v.get("ownerSub").and_then(|o| o.as_str()) != Some(self.owner_sub()) {
continue;
}
results.push(v);
}
}
}
Ok(Value::Array(results))
} else {
let mut db = self.db.lock().map_err(|e| format!("lock: {e}"))?;
let mut results = Vec::new();
let prefix = b"notes:";
let end = upper_bound(prefix);
if let Ok(entries) = db.range(prefix, &end) {
for (_, val_bytes) in entries {
if results.len() >= k {
break;
}
let s = String::from_utf8_lossy(&val_bytes).to_string();
if let Ok(v) = akurai_json::parse(&s) {
if v.get("ownerSub").and_then(|o| o.as_str()) != Some(self.owner_sub()) {
continue;
}
results.push(v);
}
}
}
Ok(Value::Array(results))
}
}
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_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 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 counter_key = "notes:_counter";
let next_id = match db
.get(counter_key.as_bytes())
.map_err(|e| format!("counter read: {e}"))?
{
Some(bytes) => {
let s = String::from_utf8_lossy(&bytes).to_string();
match s.parse::<u64>() {
Ok(id) if id < u64::MAX => id + 1,
_ => return Err("ID space exhausted".into()),
}
}
_ => 1,
};
db.insert(counter_key.as_bytes(), next_id.to_string().as_bytes())
.map_err(|e| format!("counter write: {e}"))?;
db.commit().map_err(|e| format!("commit: {e}"))?;
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 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(false)),
];
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}"))?;
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 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 Some(title) = new_title {
if let Value::Object(ref mut pairs) = note {
if let Some((_, v)) = pairs.iter_mut().find(|(k, _)| k == "title") {
*v = Value::Str(title.to_string());
}
}
}
if let Some(body) = new_body {
if let Value::Object(ref mut pairs) = note {
if let Some((_, v)) = pairs.iter_mut().find(|(k, _)| k == "body") {
*v = Value::Str(body.to_string());
}
}
}
if let Value::Object(ref mut pairs) = note {
if let Some((_, v)) = pairs.iter_mut().find(|(k, _)| k == "updatedAt") {
*v = 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}"))?;
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())
.ok_or("missing noteId")?;
let _tag_id = Self::get_arg(args, "tagId")
.and_then(|v| v.as_f64())
.ok_or("missing tagId")?;
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())
.ok_or("missing noteId")?;
let _tag_id = Self::get_arg(args, "tagId")
.and_then(|v| v.as_f64())
.ok_or("missing tagId")?;
Ok(Value::Object(vec![("ok".into(), Value::Bool(true))]))
}
}
fn upper_bound(prefix: &[u8]) -> Vec<u8> {
let mut bound = prefix.to_vec();
bound.push(0xff);
bound
}