use adk_rust::tool::SimpleToolContext;
use adk_rust::Tool;
use kyma_server::agent::{
tool_describe_table, tool_explore_schema, tool_find_references_to, tool_graph_traverse,
tool_ingest_entity, tool_link_memory_to_entity, tool_list_databases, tool_list_memories,
tool_memory_compare, tool_memory_judge, tool_memory_search, tool_memory_session_summary,
tool_recall_memory, tool_run_kql, tool_run_sql, tool_sample_rows, tool_save_memory,
tool_update_memory_importance, tool_update_memory_status,
SharedToolCtx,
};
use serde_json::{json, Value};
use std::collections::HashMap;
use std::sync::Arc;
use crate::jsonrpc::{ErrorCode, ErrorObject};
#[derive(Clone)]
pub struct ToolDispatch {
by_name: Arc<HashMap<&'static str, Arc<dyn Tool>>>,
}
impl ToolDispatch {
pub fn new(shared: SharedToolCtx) -> Self {
let mut map: HashMap<&'static str, Arc<dyn Tool>> = HashMap::with_capacity(20);
map.insert("list_databases", tool_list_databases(shared.clone()));
map.insert("describe_table", tool_describe_table(shared.clone()));
map.insert("run_sql", tool_run_sql(shared.clone()));
map.insert("run_kql", tool_run_kql(shared.clone()));
map.insert("sample_rows", tool_sample_rows(shared.clone()));
map.insert("explore_schema", tool_explore_schema(shared.clone()));
map.insert("find_references_to", tool_find_references_to(shared.clone()));
map.insert("graph_traverse", tool_graph_traverse(shared.clone()));
map.insert("memory_search", tool_memory_search(shared.clone()));
map.insert("recall_memory", tool_recall_memory(shared.clone()));
map.insert("save_memory", tool_save_memory(shared.clone()));
map.insert("list_memories", tool_list_memories(shared.clone()));
map.insert("link_memory_to_entity", tool_link_memory_to_entity(shared.clone()));
map.insert("ingest_entity", tool_ingest_entity(shared.clone()));
map.insert("update_memory_status", tool_update_memory_status(shared.clone()));
map.insert("update_memory_importance", tool_update_memory_importance(shared.clone()));
map.insert("memory_compare", tool_memory_compare(shared.clone()));
map.insert("memory_judge", tool_memory_judge(shared.clone()));
map.insert("memory_session_summary", tool_memory_session_summary(shared));
Self { by_name: Arc::new(map) }
}
pub fn list(&self) -> Vec<Value> {
let mut entries: Vec<(&'static str, Value)> = Vec::with_capacity(self.by_name.len());
for (name, tool) in self.by_name.iter() {
let input_schema = tool
.parameters_schema()
.unwrap_or_else(|| json!({"type": "object"}));
entries.push((
*name,
json!({
"name": tool.name(),
"description": tool.description(),
"inputSchema": input_schema,
}),
));
}
entries.sort_by(|a, b| a.0.cmp(b.0));
entries.into_iter().map(|(_, v)| v).collect()
}
pub async fn call(&self, name: &str, arguments: Value) -> Result<Value, ErrorObject> {
let Some(tool) = self.by_name.get(name).cloned() else {
return Err(ErrorObject::new(
ErrorCode::MethodNotFound as i64,
format!("unknown tool: {name}"),
));
};
let ctx = Arc::new(SimpleToolContext::new("kyma-mcp"));
match tool.execute(ctx, arguments).await {
Ok(value) => Ok(json!({
"content": [
{"type": "text", "text": serde_json::to_string(&value).expect("serializing serde_json::Value to string is infallible")}
],
"isError": false,
"structuredContent": value,
})),
Err(e) => Err(ErrorObject::new(
ErrorCode::InternalError as i64,
format!("tool {name}: {e}"),
)),
}
}
}