reddb_server/mcp/
server.rs

1//! MCP Server for RedDB.
2//!
3//! Runs an embedded RedDB runtime and exposes it to AI agents via the
4//! Model Context Protocol JSON-RPC transport over stdio.
5
6use crate::application::{
7    CatalogUseCases, CreateDocumentInput, CreateEdgeInput, CreateNodeInput, CreateRowInput,
8    CreateVectorInput, DeleteEntityInput, EntityUseCases, ExecuteQueryInput, GraphCentralityInput,
9    GraphClusteringInput, GraphCommunitiesInput, GraphComponentsInput, GraphCyclesInput,
10    GraphShortestPathInput, GraphTraversalInput, GraphUseCases, QueryUseCases, ScanCollectionInput,
11    SearchSimilarInput, SearchTextInput,
12};
13use crate::auth::store::AuthStore;
14use crate::auth::{AuthConfig, Role};
15use crate::json::{
16    from_str as json_from_str, to_string as json_to_string, Map, Value as JsonValue,
17};
18use crate::mcp::{protocol, tools};
19use crate::presentation::entity_json::created_entity_output_json;
20use crate::presentation::entity_json::storage_value_to_json;
21use crate::presentation::query_result_json::{runtime_query_json, runtime_stats_json};
22use crate::runtime::{
23    RedDBRuntime, RuntimeGraphCentralityAlgorithm, RuntimeGraphCommunityAlgorithm,
24    RuntimeGraphComponentsMode, RuntimeGraphDirection, RuntimeGraphPathAlgorithm,
25    RuntimeGraphTraversalStrategy,
26};
27use crate::storage::schema::Value;
28use crate::storage::EntityId;
29
30use std::io::{self, BufRead, Write};
31use std::sync::Arc;
32
33/// MCP server wrapping an embedded RedDB runtime.
34pub struct McpServer {
35    runtime: RedDBRuntime,
36    auth_store: Arc<AuthStore>,
37    initialized: bool,
38}
39
40impl McpServer {
41    /// Create a new MCP server with the given runtime.
42    pub fn new(runtime: RedDBRuntime) -> Self {
43        let auth_store = Arc::new(AuthStore::new(AuthConfig {
44            enabled: true,
45            ..Default::default()
46        }));
47        auth_store.bootstrap_from_env();
48        runtime.set_auth_store(Arc::clone(&auth_store));
49        Self {
50            runtime,
51            auth_store,
52            initialized: false,
53        }
54    }
55
56    /// Run the MCP server reading from stdin and writing to stdout.
57    ///
58    /// This blocks until stdin is closed (EOF). Diagnostic messages are
59    /// written to stderr so they do not interfere with the protocol.
60    pub fn run_stdio(&mut self) {
61        let stdin = io::stdin();
62        let stdout = io::stdout();
63        let mut reader = io::BufReader::new(stdin.lock());
64        let mut writer = io::BufWriter::new(stdout.lock());
65
66        tracing::info!(target: "reddb::mcp", "server started, waiting for messages on stdin");
67
68        loop {
69            let payload = match protocol::read_payload(&mut reader) {
70                Ok(Some(p)) => p,
71                Ok(None) => {
72                    tracing::info!(target: "reddb::mcp", "stdin closed, shutting down");
73                    break;
74                }
75                Err(e) => {
76                    tracing::error!(target: "reddb::mcp", err = %e, "read error");
77                    continue;
78                }
79            };
80
81            let request: JsonValue = match json_from_str(&payload) {
82                Ok(v) => v,
83                Err(e) => {
84                    tracing::warn!(target: "reddb::mcp", err = %e, "invalid JSON");
85                    let msg = protocol::build_error_message(None, -32700, "parse error");
86                    let _ = protocol::write_message(&mut writer, &msg);
87                    continue;
88                }
89            };
90
91            let response = self.handle_message(&request);
92            if let Some(resp) = response {
93                if let Err(e) = protocol::write_message(&mut writer, &resp) {
94                    tracing::error!(target: "reddb::mcp", err = %e, "write error");
95                    break;
96                }
97            }
98        }
99    }
100
101    /// Route a JSON-RPC message to the appropriate handler.
102    fn handle_message(&mut self, msg: &JsonValue) -> Option<String> {
103        let method = msg.get("method").and_then(|v| v.as_str())?;
104        let id = msg.get("id");
105
106        match method {
107            "initialize" => Some(self.handle_initialize(id)),
108            "initialized" | "notifications/initialized" => {
109                // Notification -- no response required.
110                None
111            }
112            "tools/list" => Some(self.handle_tools_list(id)),
113            "tools/call" => Some(self.handle_tools_call(id, msg.get("params"))),
114            "ping" => {
115                let mut result = Map::new();
116                result.insert("status".to_string(), JsonValue::String("ok".to_string()));
117                Some(protocol::build_result_message(
118                    id,
119                    JsonValue::Object(result),
120                ))
121            }
122            _ => Some(protocol::build_error_message(
123                id,
124                -32601,
125                &format!("unknown method: {}", method),
126            )),
127        }
128    }
129
130    // ------------------------------------------------------------------
131    // MCP lifecycle
132    // ------------------------------------------------------------------
133
134    fn handle_initialize(&mut self, id: Option<&JsonValue>) -> String {
135        self.initialized = true;
136
137        let mut capabilities = Map::new();
138        {
139            let mut tools_cap = Map::new();
140            tools_cap.insert("listChanged".to_string(), JsonValue::Bool(false));
141            capabilities.insert("tools".to_string(), JsonValue::Object(tools_cap));
142        }
143
144        let mut server_info = Map::new();
145        server_info.insert(
146            "name".to_string(),
147            JsonValue::String("reddb-mcp".to_string()),
148        );
149        server_info.insert(
150            "version".to_string(),
151            JsonValue::String(env!("CARGO_PKG_VERSION").to_string()),
152        );
153
154        let mut result = Map::new();
155        result.insert(
156            "protocolVersion".to_string(),
157            JsonValue::String("2024-11-05".to_string()),
158        );
159        result.insert("capabilities".to_string(), JsonValue::Object(capabilities));
160        result.insert("serverInfo".to_string(), JsonValue::Object(server_info));
161
162        protocol::build_result_message(id, JsonValue::Object(result))
163    }
164
165    // ------------------------------------------------------------------
166    // tools/list
167    // ------------------------------------------------------------------
168
169    fn handle_tools_list(&self, id: Option<&JsonValue>) -> String {
170        let defs = tools::all_tools();
171        let mut tools_json: Vec<JsonValue> = defs
172            .into_iter()
173            .map(|def| {
174                let mut obj = Map::new();
175                obj.insert("name".to_string(), JsonValue::String(def.name.to_string()));
176                obj.insert(
177                    "description".to_string(),
178                    JsonValue::String(def.description.to_string()),
179                );
180                obj.insert("inputSchema".to_string(), def.input_schema);
181                JsonValue::Object(obj)
182            })
183            .collect();
184        tools_json.push(crate::runtime::ai::mcp_ask_tool::descriptor());
185
186        let mut result = Map::new();
187        result.insert("tools".to_string(), JsonValue::Array(tools_json));
188        protocol::build_result_message(id, JsonValue::Object(result))
189    }
190
191    // ------------------------------------------------------------------
192    // tools/call dispatcher
193    // ------------------------------------------------------------------
194
195    fn handle_tools_call(&self, id: Option<&JsonValue>, params: Option<&JsonValue>) -> String {
196        let name = params.and_then(|p| p.get("name")).and_then(|v| v.as_str());
197        let name = match name {
198            Some(n) => n,
199            None => {
200                return protocol::build_error_message(id, -32602, "missing tool name");
201            }
202        };
203
204        let empty = JsonValue::Object(Map::new());
205        let args = params.and_then(|p| p.get("arguments")).unwrap_or(&empty);
206
207        let result = match name {
208            "reddb_query" => self.tool_query(args),
209            "reddb_collections" => self.tool_collections(),
210            "reddb_insert_row" => self.tool_insert_row(args),
211            "reddb_insert_node" => self.tool_insert_node(args),
212            "reddb_insert_edge" => self.tool_insert_edge(args),
213            "reddb_insert_vector" => self.tool_insert_vector(args),
214            "reddb_insert_document" => self.tool_insert_document(args),
215            "reddb_kv_get" => self.tool_kv_get(args),
216            "reddb_kv_set" => self.tool_kv_set(args),
217            "reddb_kv_invalidate_tags" => self.tool_kv_invalidate_tags(args),
218            "reddb_config_get" => self.tool_config_get(args),
219            "reddb_config_put" => self.tool_config_put(args),
220            "reddb_config_resolve" => self.tool_config_resolve(args),
221            "reddb_vault_get" => self.tool_vault_get(args),
222            "reddb_vault_put" => self.tool_vault_put(args),
223            "reddb_vault_unseal" => self.tool_vault_unseal(args),
224            "reddb_delete" => self.tool_delete(args),
225            "reddb_search_vector" => self.tool_search_vector(args),
226            "reddb_search_text" => self.tool_search_text(args),
227            "reddb_health" => self.tool_health(),
228            "reddb_graph_traverse" => self.tool_graph_traverse(args),
229            "reddb_graph_shortest_path" => self.tool_graph_shortest_path(args),
230            "reddb_update" => self.tool_update(args),
231            "reddb_scan" => self.tool_scan(args),
232            "reddb_graph_centrality" => self.tool_graph_centrality(args),
233            "reddb_graph_community" => self.tool_graph_community(args),
234            "reddb_graph_components" => self.tool_graph_components(args),
235            "reddb_graph_cycles" => self.tool_graph_cycles(args),
236            "reddb_graph_clustering" => self.tool_graph_clustering(args),
237            "reddb_create_collection" => self.tool_create_collection(args),
238            "reddb_drop_collection" => self.tool_drop_collection(args),
239            "reddb_auth_bootstrap" => self.tool_auth_bootstrap(args),
240            "reddb_auth_create_user" => self.tool_auth_create_user(args),
241            "reddb_auth_login" => self.tool_auth_login(args),
242            "reddb_auth_create_api_key" => self.tool_auth_create_api_key(args),
243            "reddb_auth_list_users" => self.tool_auth_list_users(),
244            crate::runtime::ai::mcp_ask_tool::TOOL_NAME => self.tool_ask(args),
245            _ => Err(format!("unknown tool: {name}")),
246        };
247
248        match result {
249            Ok(text) => {
250                let mut content = Map::new();
251                content.insert("type".to_string(), JsonValue::String("text".to_string()));
252                content.insert("text".to_string(), JsonValue::String(text));
253
254                let mut result_obj = Map::new();
255                result_obj.insert(
256                    "content".to_string(),
257                    JsonValue::Array(vec![JsonValue::Object(content)]),
258                );
259                protocol::build_result_message(id, JsonValue::Object(result_obj))
260            }
261            Err(err) => {
262                let mut content = Map::new();
263                content.insert("type".to_string(), JsonValue::String("text".to_string()));
264                content.insert("text".to_string(), JsonValue::String(err.clone()));
265
266                let mut result_obj = Map::new();
267                result_obj.insert(
268                    "content".to_string(),
269                    JsonValue::Array(vec![JsonValue::Object(content)]),
270                );
271                result_obj.insert("isError".to_string(), JsonValue::Bool(true));
272                protocol::build_result_message(id, JsonValue::Object(result_obj))
273            }
274        }
275    }
276
277    // ------------------------------------------------------------------
278    // Tool implementations
279    // ------------------------------------------------------------------
280
281    fn tool_query(&self, args: &JsonValue) -> Result<String, String> {
282        let sql = args
283            .get("sql")
284            .and_then(|v| v.as_str())
285            .ok_or("missing required field 'sql'")?;
286
287        // Optional positional `$N` bind parameters. Decoded via the same
288        // helper the JSON-RPC stdio path uses (#358), so MCP, embedded
289        // stdio, and HTTP all bind via one codec.
290        if let Some(raw_params) = args.get("params") {
291            let arr = raw_params
292                .as_array()
293                .ok_or_else(|| "'params' must be an array".to_string())?;
294            let binds: Vec<Value> = arr
295                .iter()
296                .map(crate::rpc_stdio::json_value_to_schema_value)
297                .collect();
298
299            use crate::storage::query::modes::parse_multi;
300            use crate::storage::query::user_params;
301            let parsed = parse_multi(sql).map_err(|e| format!("{}", e))?;
302            let bound = user_params::bind(&parsed, &binds).map_err(|e| format!("{}", e))?;
303            let result = self
304                .runtime
305                .execute_query_expr(bound)
306                .map_err(|e| format!("{}", e))?;
307            let json = runtime_query_json(&result, &None, &None);
308            return json_to_string(&json).map_err(|e| format!("serialization error: {}", e));
309        }
310
311        let uc = QueryUseCases::new(&self.runtime);
312        let result = uc
313            .execute(ExecuteQueryInput {
314                query: sql.to_string(),
315            })
316            .map_err(|e| format!("{}", e))?;
317
318        let json = runtime_query_json(&result, &None, &None);
319        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
320    }
321
322    fn tool_ask(&self, args: &JsonValue) -> Result<String, String> {
323        let invocation =
324            crate::runtime::ai::mcp_ask_tool::parse(args).map_err(format_mcp_ask_parse_error)?;
325        let ask = crate::storage::query::ast::AskQuery {
326            explain: false,
327            question: invocation.question,
328            question_param: None,
329            provider: invocation.using,
330            model: invocation.model,
331            depth: invocation.depth.map(|v| v as usize),
332            limit: invocation.limit.map(|v| v as usize),
333            min_score: invocation.min_score.map(|v| v as f32),
334            collection: None,
335            temperature: invocation.temperature.map(|v| v as f32),
336            seed: invocation.seed,
337            strict: invocation.strict.unwrap_or(true),
338            stream: false,
339            cache: if matches!(invocation.nocache, Some(true)) {
340                crate::storage::query::ast::AskCacheClause::NoCache
341            } else if let Some(ttl) = invocation.cache_ttl {
342                crate::storage::query::ast::AskCacheClause::CacheTtl(ttl)
343            } else {
344                crate::storage::query::ast::AskCacheClause::Default
345            },
346        };
347        let result = self
348            .runtime
349            .execute_ask("ASK <mcp>", &ask)
350            .map_err(|e| format!("{}", e))?;
351        let json = crate::rpc_stdio::query_result_to_json(&result);
352        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
353    }
354
355    fn tool_collections(&self) -> Result<String, String> {
356        let uc = CatalogUseCases::new(&self.runtime);
357        let collections = uc.collections();
358        let json = JsonValue::Array(collections.into_iter().map(JsonValue::String).collect());
359        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
360    }
361
362    fn tool_insert_row(&self, args: &JsonValue) -> Result<String, String> {
363        let collection = args
364            .get("collection")
365            .and_then(|v| v.as_str())
366            .ok_or("missing required field 'collection'")?;
367        let data = args
368            .get("data")
369            .and_then(|v| v.as_object())
370            .ok_or("missing required field 'data' (must be an object)")?;
371
372        let mut fields = Vec::new();
373        for (key, value) in data {
374            let sv = crate::application::entity::json_to_storage_value(value)
375                .map_err(|e| format!("{}", e))?;
376            fields.push((key.clone(), sv));
377        }
378
379        let metadata = parse_metadata_arg(args)?;
380
381        let uc = EntityUseCases::new(&self.runtime);
382        let output = uc
383            .create_row(CreateRowInput {
384                collection: collection.to_string(),
385                fields,
386                metadata,
387                node_links: vec![],
388                vector_links: vec![],
389            })
390            .map_err(|e| format!("{}", e))?;
391
392        let json = created_entity_output_json(&output);
393        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
394    }
395
396    fn tool_insert_node(&self, args: &JsonValue) -> Result<String, String> {
397        let collection = args
398            .get("collection")
399            .and_then(|v| v.as_str())
400            .ok_or("missing required field 'collection'")?;
401        let label = args
402            .get("label")
403            .and_then(|v| v.as_str())
404            .ok_or("missing required field 'label'")?;
405        let node_type = args
406            .get("node_type")
407            .and_then(|v| v.as_str())
408            .map(String::from);
409
410        let mut properties = Vec::new();
411        if let Some(props) = args.get("properties").and_then(|v| v.as_object()) {
412            for (key, value) in props {
413                let sv = crate::application::entity::json_to_storage_value(value)
414                    .map_err(|e| format!("{}", e))?;
415                properties.push((key.clone(), sv));
416            }
417        }
418
419        let metadata = parse_metadata_arg(args)?;
420
421        let uc = EntityUseCases::new(&self.runtime);
422        let output = uc
423            .create_node(CreateNodeInput {
424                collection: collection.to_string(),
425                label: label.to_string(),
426                node_type,
427                properties,
428                metadata,
429                embeddings: vec![],
430                table_links: vec![],
431                node_links: vec![],
432            })
433            .map_err(|e| format!("{}", e))?;
434
435        let json = created_entity_output_json(&output);
436        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
437    }
438
439    fn tool_insert_edge(&self, args: &JsonValue) -> Result<String, String> {
440        let collection = args
441            .get("collection")
442            .and_then(|v| v.as_str())
443            .ok_or("missing required field 'collection'")?;
444        let label = args
445            .get("label")
446            .and_then(|v| v.as_str())
447            .ok_or("missing required field 'label'")?;
448        let from_id = args
449            .get("from")
450            .and_then(|v| v.as_u64())
451            .ok_or("missing required field 'from' (integer)")?;
452        let to_id = args
453            .get("to")
454            .and_then(|v| v.as_u64())
455            .ok_or("missing required field 'to' (integer)")?;
456        let weight = args
457            .get("weight")
458            .and_then(|v| v.as_f64())
459            .map(|w| w as f32);
460
461        let mut properties = Vec::new();
462        if let Some(props) = args.get("properties").and_then(|v| v.as_object()) {
463            for (key, value) in props {
464                let sv = crate::application::entity::json_to_storage_value(value)
465                    .map_err(|e| format!("{}", e))?;
466                properties.push((key.clone(), sv));
467            }
468        }
469
470        let metadata = parse_metadata_arg(args)?;
471
472        let uc = EntityUseCases::new(&self.runtime);
473        let output = uc
474            .create_edge(CreateEdgeInput {
475                collection: collection.to_string(),
476                label: label.to_string(),
477                from: EntityId::new(from_id),
478                to: EntityId::new(to_id),
479                weight,
480                properties,
481                metadata,
482            })
483            .map_err(|e| format!("{}", e))?;
484
485        let json = created_entity_output_json(&output);
486        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
487    }
488
489    fn tool_insert_vector(&self, args: &JsonValue) -> Result<String, String> {
490        let collection = args
491            .get("collection")
492            .and_then(|v| v.as_str())
493            .ok_or("missing required field 'collection'")?;
494        let dense_arr = args
495            .get("dense")
496            .and_then(|v| v.as_array())
497            .ok_or("missing required field 'dense' (array of numbers)")?;
498
499        let mut dense = Vec::with_capacity(dense_arr.len());
500        for v in dense_arr {
501            dense.push(
502                v.as_f64()
503                    .ok_or("'dense' array must contain only numbers")? as f32,
504            );
505        }
506        if dense.is_empty() {
507            return Err("'dense' vector cannot be empty".to_string());
508        }
509
510        let content = args
511            .get("content")
512            .and_then(|v| v.as_str())
513            .map(String::from);
514        let metadata = parse_metadata_arg(args)?;
515
516        let uc = EntityUseCases::new(&self.runtime);
517        let output = uc
518            .create_vector(CreateVectorInput {
519                collection: collection.to_string(),
520                dense,
521                content,
522                metadata,
523                link_row: None,
524                link_node: None,
525            })
526            .map_err(|e| format!("{}", e))?;
527
528        let json = created_entity_output_json(&output);
529        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
530    }
531
532    fn tool_insert_document(&self, args: &JsonValue) -> Result<String, String> {
533        let collection = args
534            .get("collection")
535            .and_then(|v| v.as_str())
536            .ok_or("missing required field 'collection'")?;
537        let body = args.get("body").ok_or("missing required field 'body'")?;
538
539        let metadata = parse_metadata_arg(args)?;
540
541        let uc = EntityUseCases::new(&self.runtime);
542        let output = uc
543            .create_document(CreateDocumentInput {
544                collection: collection.to_string(),
545                body: body.clone(),
546                metadata,
547                node_links: vec![],
548                vector_links: vec![],
549            })
550            .map_err(|e| format!("{}", e))?;
551
552        let json = created_entity_output_json(&output);
553        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
554    }
555
556    fn tool_kv_get(&self, args: &JsonValue) -> Result<String, String> {
557        let collection = args
558            .get("collection")
559            .and_then(|v| v.as_str())
560            .ok_or("missing required field 'collection'")?;
561        let key = args
562            .get("key")
563            .and_then(|v| v.as_str())
564            .ok_or("missing required field 'key'")?;
565
566        let uc = EntityUseCases::new(&self.runtime);
567        match uc.get_kv(collection, key).map_err(|e| format!("{}", e))? {
568            Some((value, entity_id)) => {
569                let mut obj = Map::new();
570                obj.insert("found".to_string(), JsonValue::Bool(true));
571                obj.insert("key".to_string(), JsonValue::String(key.to_string()));
572                obj.insert("value".to_string(), storage_value_to_json(&value));
573                obj.insert("rid".to_string(), JsonValue::Number(entity_id.raw() as f64));
574                obj.insert("kind".to_string(), JsonValue::String("kv".to_string()));
575                json_to_string(&JsonValue::Object(obj))
576                    .map_err(|e| format!("serialization error: {}", e))
577            }
578            None => {
579                let mut obj = Map::new();
580                obj.insert("found".to_string(), JsonValue::Bool(false));
581                obj.insert("key".to_string(), JsonValue::String(key.to_string()));
582                json_to_string(&JsonValue::Object(obj))
583                    .map_err(|e| format!("serialization error: {}", e))
584            }
585        }
586    }
587
588    fn tool_kv_set(&self, args: &JsonValue) -> Result<String, String> {
589        let collection = args
590            .get("collection")
591            .and_then(|v| v.as_str())
592            .ok_or("missing required field 'collection'")?;
593        let key = args
594            .get("key")
595            .and_then(|v| v.as_str())
596            .ok_or("missing required field 'key'")?;
597        let value_arg = args.get("value").ok_or("missing required field 'value'")?;
598
599        let sv = crate::application::entity::json_to_storage_value(value_arg)
600            .map_err(|e| format!("{}", e))?;
601
602        let metadata = parse_metadata_arg(args)?;
603
604        let tags = parse_string_array_arg(args, "tags")?;
605        let ops = crate::runtime::impl_kv::KvAtomicOps::new(&self.runtime);
606        let (_, id) = ops
607            .set_with_tags_and_metadata(collection, key, sv, None, &tags, false, metadata)
608            .map_err(|e| format!("{}", e))?;
609
610        let mut obj = Map::new();
611        obj.insert("ok".to_string(), JsonValue::Bool(true));
612        obj.insert("rid".to_string(), JsonValue::Number(id.raw() as f64));
613        obj.insert("kind".to_string(), JsonValue::String("kv".to_string()));
614        obj.insert(
615            "tags".to_string(),
616            JsonValue::Array(tags.into_iter().map(JsonValue::String).collect()),
617        );
618        json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
619    }
620
621    fn tool_kv_invalidate_tags(&self, args: &JsonValue) -> Result<String, String> {
622        let collection = args
623            .get("collection")
624            .and_then(|v| v.as_str())
625            .ok_or("missing required field 'collection'")?;
626        let tags = parse_string_array_arg(args, "tags")?;
627        if tags.is_empty() {
628            return Err("missing required field 'tags'".to_string());
629        }
630        let ops = crate::runtime::impl_kv::KvAtomicOps::new(&self.runtime);
631        let count = ops
632            .invalidate_tags(collection, &tags)
633            .map_err(|e| format!("{}", e))?;
634
635        let mut obj = Map::new();
636        obj.insert("ok".to_string(), JsonValue::Bool(true));
637        obj.insert("invalidated".to_string(), JsonValue::Number(count as f64));
638        obj.insert(
639            "tags".to_string(),
640            JsonValue::Array(tags.into_iter().map(JsonValue::String).collect()),
641        );
642        json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
643    }
644
645    fn tool_config_get(&self, args: &JsonValue) -> Result<String, String> {
646        let collection = mcp_keyed_ident(get_str_field(args, "collection")?)?;
647        let key = mcp_keyed_ident(get_str_field(args, "key")?)?;
648        self.tool_keyed_query(format!("GET CONFIG {collection} {key}"))
649    }
650
651    fn tool_config_put(&self, args: &JsonValue) -> Result<String, String> {
652        reject_mcp_volatile_options(args, "CONFIG")?;
653        let collection = mcp_keyed_ident(get_str_field(args, "collection")?)?;
654        let key = mcp_keyed_ident(get_str_field(args, "key")?)?;
655        let tags = parse_string_array_arg(args, "tags")?;
656        let literal = if let Some(secret_ref) = args.get("secret_ref") {
657            let object = secret_ref
658                .as_object()
659                .ok_or("field 'secret_ref' must be an object")?;
660            let ref_collection = object
661                .get("collection")
662                .and_then(|v| v.as_str())
663                .ok_or_else(|| "secret_ref.collection is required".to_string())
664                .and_then(mcp_keyed_ident)?;
665            let ref_key = object
666                .get("key")
667                .and_then(|v| v.as_str())
668                .ok_or_else(|| "secret_ref.key is required".to_string())
669                .and_then(mcp_keyed_ident)?;
670            format!("SECRET_REF(vault, {ref_collection}.{ref_key})")
671        } else {
672            mcp_value_literal(args.get("value").ok_or("missing required field 'value'")?)?
673        };
674        let mut sql = format!("PUT CONFIG {collection} {key} = {literal}");
675        append_mcp_tags_clause(&mut sql, &tags);
676        self.tool_keyed_query(sql)
677    }
678
679    fn tool_config_resolve(&self, args: &JsonValue) -> Result<String, String> {
680        let collection = mcp_keyed_ident(get_str_field(args, "collection")?)?;
681        let key = mcp_keyed_ident(get_str_field(args, "key")?)?;
682        self.tool_keyed_query(format!("RESOLVE CONFIG {collection} {key}"))
683    }
684
685    fn tool_vault_get(&self, args: &JsonValue) -> Result<String, String> {
686        let collection = mcp_keyed_ident(get_str_field(args, "collection")?)?;
687        let key = mcp_keyed_ident(get_str_field(args, "key")?)?;
688        self.tool_keyed_query(format!("VAULT GET {collection}.{key}"))
689    }
690
691    fn tool_vault_put(&self, args: &JsonValue) -> Result<String, String> {
692        reject_mcp_volatile_options(args, "VAULT")?;
693        let collection = mcp_keyed_ident(get_str_field(args, "collection")?)?;
694        let key = mcp_keyed_ident(get_str_field(args, "key")?)?;
695        let value = mcp_value_literal(args.get("value").ok_or("missing required field 'value'")?)?;
696        let tags = parse_string_array_arg(args, "tags")?;
697        let mut sql = format!("VAULT PUT {collection}.{key} = {value}");
698        append_mcp_tags_clause(&mut sql, &tags);
699        self.tool_keyed_query(sql)
700    }
701
702    fn tool_vault_unseal(&self, args: &JsonValue) -> Result<String, String> {
703        let collection = mcp_keyed_ident(get_str_field(args, "collection")?)?;
704        let key = mcp_keyed_ident(get_str_field(args, "key")?)?;
705        self.tool_keyed_query(format!("UNSEAL VAULT {collection}.{key}"))
706    }
707
708    fn tool_keyed_query(&self, sql: String) -> Result<String, String> {
709        let uc = QueryUseCases::new(&self.runtime);
710        let result = uc
711            .execute(ExecuteQueryInput { query: sql })
712            .map_err(|e| format!("{}", e))?;
713        let json = runtime_query_json(&result, &None, &None);
714        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
715    }
716
717    fn tool_delete(&self, args: &JsonValue) -> Result<String, String> {
718        let collection = args
719            .get("collection")
720            .and_then(|v| v.as_str())
721            .ok_or("missing required field 'collection'")?;
722        let id = args
723            .get("id")
724            .and_then(|v| v.as_u64())
725            .ok_or("missing required field 'id' (integer)")?;
726
727        let uc = EntityUseCases::new(&self.runtime);
728        let output = uc
729            .delete(DeleteEntityInput {
730                collection: collection.to_string(),
731                id: EntityId::new(id),
732            })
733            .map_err(|e| format!("{}", e))?;
734
735        let mut obj = Map::new();
736        obj.insert("deleted".to_string(), JsonValue::Bool(output.deleted));
737        obj.insert("id".to_string(), JsonValue::Number(output.id.raw() as f64));
738        json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
739    }
740
741    fn tool_search_vector(&self, args: &JsonValue) -> Result<String, String> {
742        let collection = args
743            .get("collection")
744            .and_then(|v| v.as_str())
745            .ok_or("missing required field 'collection'")?;
746        let vector_arr = args
747            .get("vector")
748            .and_then(|v| v.as_array())
749            .ok_or("missing required field 'vector' (array of numbers)")?;
750
751        let mut vector = Vec::with_capacity(vector_arr.len());
752        for v in vector_arr {
753            vector.push(
754                v.as_f64()
755                    .ok_or("'vector' array must contain only numbers")? as f32,
756            );
757        }
758        let k = args
759            .get("k")
760            .and_then(|v| v.as_u64())
761            .map(|v| v as usize)
762            .unwrap_or(10);
763        let min_score = args
764            .get("min_score")
765            .and_then(|v| v.as_f64())
766            .map(|v| v as f32)
767            .unwrap_or(0.0);
768
769        let uc = QueryUseCases::new(&self.runtime);
770        let results = uc
771            .search_similar(SearchSimilarInput {
772                collection: collection.to_string(),
773                vector,
774                k,
775                min_score,
776                text: None,
777                provider: None,
778            })
779            .map_err(|e| format!("{}", e))?;
780
781        let items: Vec<JsonValue> = results
782            .iter()
783            .map(|r| {
784                let mut obj = Map::new();
785                obj.insert(
786                    "rid".to_string(),
787                    JsonValue::Number(r.entity_id.raw() as f64),
788                );
789                obj.insert("kind".to_string(), JsonValue::String("vector".to_string()));
790                obj.insert("score".to_string(), JsonValue::Number(r.score as f64));
791                obj.insert("distance".to_string(), JsonValue::Number(r.distance as f64));
792                JsonValue::Object(obj)
793            })
794            .collect();
795
796        let mut obj = Map::new();
797        obj.insert("count".to_string(), JsonValue::Number(items.len() as f64));
798        obj.insert("results".to_string(), JsonValue::Array(items));
799        json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
800    }
801
802    fn tool_search_text(&self, args: &JsonValue) -> Result<String, String> {
803        let query = args
804            .get("query")
805            .and_then(|v| v.as_str())
806            .ok_or("missing required field 'query'")?;
807
808        let collections = args
809            .get("collections")
810            .and_then(|v| v.as_array())
811            .map(|arr| {
812                arr.iter()
813                    .filter_map(|v| v.as_str().map(String::from))
814                    .collect::<Vec<_>>()
815            });
816        let limit = args
817            .get("limit")
818            .and_then(|v| v.as_u64())
819            .map(|v| v as usize);
820        let fuzzy = args.get("fuzzy").and_then(|v| v.as_bool()).unwrap_or(false);
821
822        let uc = QueryUseCases::new(&self.runtime);
823        let result = uc
824            .search_text(SearchTextInput {
825                query: query.to_string(),
826                collections,
827                entity_types: None,
828                capabilities: None,
829                fields: None,
830                limit,
831                fuzzy,
832            })
833            .map_err(|e| format!("{}", e))?;
834
835        let items: Vec<JsonValue> = result
836            .matches
837            .iter()
838            .map(|m| {
839                let mut obj = Map::new();
840                obj.insert(
841                    "rid".to_string(),
842                    JsonValue::Number(m.entity.id.raw() as f64),
843                );
844                obj.insert(
845                    "kind".to_string(),
846                    JsonValue::String(m.entity.kind.storage_type().to_string()),
847                );
848                obj.insert("score".to_string(), JsonValue::Number(m.score as f64));
849                JsonValue::Object(obj)
850            })
851            .collect();
852
853        let mut obj = Map::new();
854        obj.insert("count".to_string(), JsonValue::Number(items.len() as f64));
855        obj.insert("results".to_string(), JsonValue::Array(items));
856        json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
857    }
858
859    fn tool_health(&self) -> Result<String, String> {
860        let uc = CatalogUseCases::new(&self.runtime);
861        let stats = uc.stats();
862        let json = runtime_stats_json(&stats);
863        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
864    }
865
866    fn tool_graph_traverse(&self, args: &JsonValue) -> Result<String, String> {
867        let source = args
868            .get("source")
869            .and_then(|v| v.as_str())
870            .ok_or("missing required field 'source'")?;
871        let direction = parse_direction(args.get("direction").and_then(|v| v.as_str()));
872        let max_depth = args
873            .get("max_depth")
874            .and_then(|v| v.as_u64())
875            .map(|v| v as usize)
876            .unwrap_or(3);
877        let strategy = match args.get("strategy").and_then(|v| v.as_str()) {
878            Some("dfs") => RuntimeGraphTraversalStrategy::Dfs,
879            _ => RuntimeGraphTraversalStrategy::Bfs,
880        };
881
882        let uc = GraphUseCases::new(&self.runtime);
883        let result = uc
884            .traverse(GraphTraversalInput {
885                source: source.to_string(),
886                direction,
887                max_depth,
888                strategy,
889                edge_labels: None,
890                projection: None,
891            })
892            .map_err(|e| format!("{}", e))?;
893
894        let visits: Vec<JsonValue> = result
895            .visits
896            .iter()
897            .map(|v| {
898                let mut obj = Map::new();
899                obj.insert("depth".to_string(), JsonValue::Number(v.depth as f64));
900                obj.insert("node_id".to_string(), JsonValue::String(v.node.id.clone()));
901                obj.insert("label".to_string(), JsonValue::String(v.node.label.clone()));
902                obj.insert(
903                    "node_type".to_string(),
904                    JsonValue::String(v.node.node_type.clone()),
905                );
906                JsonValue::Object(obj)
907            })
908            .collect();
909
910        let edges: Vec<JsonValue> = result
911            .edges
912            .iter()
913            .map(|e| {
914                let mut obj = Map::new();
915                obj.insert("source".to_string(), JsonValue::String(e.source.clone()));
916                obj.insert("target".to_string(), JsonValue::String(e.target.clone()));
917                obj.insert(
918                    "edge_type".to_string(),
919                    JsonValue::String(e.edge_type.clone()),
920                );
921                obj.insert("weight".to_string(), JsonValue::Number(e.weight as f64));
922                JsonValue::Object(obj)
923            })
924            .collect();
925
926        let mut obj = Map::new();
927        obj.insert(
928            "source".to_string(),
929            JsonValue::String(result.source.clone()),
930        );
931        obj.insert(
932            "visit_count".to_string(),
933            JsonValue::Number(visits.len() as f64),
934        );
935        obj.insert("visits".to_string(), JsonValue::Array(visits));
936        obj.insert("edges".to_string(), JsonValue::Array(edges));
937        json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
938    }
939
940    fn tool_graph_shortest_path(&self, args: &JsonValue) -> Result<String, String> {
941        let source = args
942            .get("source")
943            .and_then(|v| v.as_str())
944            .ok_or("missing required field 'source'")?;
945        let target = args
946            .get("target")
947            .and_then(|v| v.as_str())
948            .ok_or("missing required field 'target'")?;
949        let direction = parse_direction(args.get("direction").and_then(|v| v.as_str()));
950        let algorithm = match args.get("algorithm").and_then(|v| v.as_str()) {
951            Some("astar") | Some("a*") => RuntimeGraphPathAlgorithm::AStar,
952            Some("bellman_ford") | Some("bellmanford") => RuntimeGraphPathAlgorithm::BellmanFord,
953            Some("dijkstra") => RuntimeGraphPathAlgorithm::Dijkstra,
954            _ => RuntimeGraphPathAlgorithm::Bfs,
955        };
956
957        let uc = GraphUseCases::new(&self.runtime);
958        let result = uc
959            .shortest_path(GraphShortestPathInput {
960                source: source.to_string(),
961                target: target.to_string(),
962                direction,
963                algorithm,
964                edge_labels: None,
965                projection: None,
966            })
967            .map_err(|e| format!("{}", e))?;
968
969        let mut obj = Map::new();
970        obj.insert(
971            "source".to_string(),
972            JsonValue::String(result.source.clone()),
973        );
974        obj.insert(
975            "target".to_string(),
976            JsonValue::String(result.target.clone()),
977        );
978        obj.insert(
979            "nodes_visited".to_string(),
980            JsonValue::Number(result.nodes_visited as f64),
981        );
982
983        match &result.path {
984            Some(path) => {
985                obj.insert("found".to_string(), JsonValue::Bool(true));
986                obj.insert(
987                    "hop_count".to_string(),
988                    JsonValue::Number(path.hop_count as f64),
989                );
990                obj.insert(
991                    "total_weight".to_string(),
992                    JsonValue::Number(path.total_weight),
993                );
994                let nodes_json: Vec<JsonValue> = path
995                    .nodes
996                    .iter()
997                    .map(|n| {
998                        let mut nobj = Map::new();
999                        nobj.insert("id".to_string(), JsonValue::String(n.id.clone()));
1000                        nobj.insert("label".to_string(), JsonValue::String(n.label.clone()));
1001                        JsonValue::Object(nobj)
1002                    })
1003                    .collect();
1004                obj.insert("nodes".to_string(), JsonValue::Array(nodes_json));
1005            }
1006            None => {
1007                obj.insert("found".to_string(), JsonValue::Bool(false));
1008            }
1009        }
1010
1011        json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
1012    }
1013
1014    fn tool_update(&self, args: &JsonValue) -> Result<String, String> {
1015        let collection = get_str_field(args, "collection")?;
1016        let set_obj = args.get("set").ok_or("missing 'set'")?;
1017        let where_clause = args
1018            .get("where_filter")
1019            .and_then(|v| v.as_str())
1020            .unwrap_or("");
1021
1022        // Build UPDATE SQL and execute via runtime
1023        let mut sql = format!("UPDATE {} SET ", collection);
1024        if let Some(obj) = set_obj.as_object() {
1025            let assignments: Vec<String> = obj
1026                .iter()
1027                .map(|(k, v)| {
1028                    let val_str = match v {
1029                        JsonValue::String(s) => format!("'{}'", s),
1030                        JsonValue::Number(n) => n.to_string(),
1031                        JsonValue::Bool(b) => b.to_string(),
1032                        _ => format!("'{}'", v),
1033                    };
1034                    format!("{} = {}", k, val_str)
1035                })
1036                .collect();
1037            sql.push_str(&assignments.join(", "));
1038        } else {
1039            return Err("'set' must be a JSON object".to_string());
1040        }
1041        if !where_clause.is_empty() {
1042            sql.push_str(&format!(" WHERE {}", where_clause));
1043        }
1044
1045        let uc = QueryUseCases::new(&self.runtime);
1046        let result = uc
1047            .execute(ExecuteQueryInput { query: sql })
1048            .map_err(|e| format!("{}", e))?;
1049
1050        let mut resp = Map::new();
1051        resp.insert("ok".into(), JsonValue::Bool(true));
1052        resp.insert(
1053            "affected_rows".into(),
1054            JsonValue::Number(result.affected_rows as f64),
1055        );
1056        json_to_string(&JsonValue::Object(resp)).map_err(|e| format!("serialization error: {}", e))
1057    }
1058
1059    fn tool_scan(&self, args: &JsonValue) -> Result<String, String> {
1060        let collection = get_str_field(args, "collection")?;
1061        let limit = args
1062            .get("limit")
1063            .and_then(|v| v.as_u64())
1064            .map(|v| v as usize)
1065            .unwrap_or(10);
1066        let offset = args
1067            .get("offset")
1068            .and_then(|v| v.as_u64())
1069            .map(|v| v as usize)
1070            .unwrap_or(0);
1071
1072        let uc = QueryUseCases::new(&self.runtime);
1073        let page = uc
1074            .scan(ScanCollectionInput {
1075                collection: collection.to_string(),
1076                offset,
1077                limit,
1078            })
1079            .map_err(|e| format!("{}", e))?;
1080
1081        let json = crate::presentation::entity_json::scan_page_json(&page);
1082        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
1083    }
1084
1085    fn tool_graph_centrality(&self, args: &JsonValue) -> Result<String, String> {
1086        let algorithm_str = get_str_field(args, "algorithm")?;
1087        let algo = match algorithm_str {
1088            "degree" => RuntimeGraphCentralityAlgorithm::Degree,
1089            "closeness" => RuntimeGraphCentralityAlgorithm::Closeness,
1090            "betweenness" => RuntimeGraphCentralityAlgorithm::Betweenness,
1091            "eigenvector" => RuntimeGraphCentralityAlgorithm::Eigenvector,
1092            "pagerank" => RuntimeGraphCentralityAlgorithm::PageRank,
1093            _ => return Err(format!("unknown algorithm: {algorithm_str}")),
1094        };
1095
1096        let uc = GraphUseCases::new(&self.runtime);
1097        let result = uc
1098            .centrality(GraphCentralityInput {
1099                algorithm: algo,
1100                top_k: 100,
1101                normalize: true,
1102                max_iterations: None,
1103                epsilon: None,
1104                alpha: None,
1105                projection: None,
1106            })
1107            .map_err(|e| format!("{}", e))?;
1108
1109        let json = crate::presentation::graph_json::graph_centrality_json(&result);
1110        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
1111    }
1112
1113    fn tool_graph_community(&self, args: &JsonValue) -> Result<String, String> {
1114        let algorithm_str = get_str_field(args, "algorithm")?;
1115        let algo = match algorithm_str {
1116            "label_propagation" => RuntimeGraphCommunityAlgorithm::LabelPropagation,
1117            "louvain" => RuntimeGraphCommunityAlgorithm::Louvain,
1118            _ => return Err(format!("unknown algorithm: {algorithm_str}")),
1119        };
1120        let max_iterations = args
1121            .get("max_iterations")
1122            .and_then(|v| v.as_u64())
1123            .map(|v| v as usize);
1124
1125        let uc = GraphUseCases::new(&self.runtime);
1126        let result = uc
1127            .communities(GraphCommunitiesInput {
1128                algorithm: algo,
1129                min_size: 1,
1130                max_iterations,
1131                resolution: None,
1132                projection: None,
1133            })
1134            .map_err(|e| format!("{}", e))?;
1135
1136        let json = crate::presentation::graph_json::graph_community_json(&result);
1137        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
1138    }
1139
1140    fn tool_graph_components(&self, args: &JsonValue) -> Result<String, String> {
1141        let mode = match args.get("mode").and_then(|v| v.as_str()) {
1142            Some("strongly_connected") => RuntimeGraphComponentsMode::Strong,
1143            _ => RuntimeGraphComponentsMode::Weak,
1144        };
1145
1146        let uc = GraphUseCases::new(&self.runtime);
1147        let result = uc
1148            .components(GraphComponentsInput {
1149                mode,
1150                min_size: 1,
1151                projection: None,
1152            })
1153            .map_err(|e| format!("{}", e))?;
1154
1155        let json = crate::presentation::graph_json::graph_components_json(&result);
1156        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
1157    }
1158
1159    fn tool_graph_cycles(&self, args: &JsonValue) -> Result<String, String> {
1160        let max_length = args
1161            .get("max_length")
1162            .and_then(|v| v.as_u64())
1163            .map(|v| v as usize)
1164            .unwrap_or(10);
1165        let max_cycles = args
1166            .get("max_cycles")
1167            .and_then(|v| v.as_u64())
1168            .map(|v| v as usize)
1169            .unwrap_or(100);
1170
1171        let uc = GraphUseCases::new(&self.runtime);
1172        let result = uc
1173            .cycles(GraphCyclesInput {
1174                max_length,
1175                max_cycles,
1176                projection: None,
1177            })
1178            .map_err(|e| format!("{}", e))?;
1179
1180        let json = crate::presentation::graph_json::graph_cycles_json(&result);
1181        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
1182    }
1183
1184    fn tool_graph_clustering(&self, _args: &JsonValue) -> Result<String, String> {
1185        let uc = GraphUseCases::new(&self.runtime);
1186        let result = uc
1187            .clustering(GraphClusteringInput {
1188                top_k: 100,
1189                include_triangles: true,
1190                projection: None,
1191            })
1192            .map_err(|e| format!("{}", e))?;
1193
1194        let json = crate::presentation::graph_json::graph_clustering_json(&result);
1195        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
1196    }
1197
1198    fn tool_create_collection(&self, args: &JsonValue) -> Result<String, String> {
1199        let name = get_str_field(args, "name")?;
1200        self.runtime
1201            .db()
1202            .store()
1203            .create_collection(name)
1204            .map_err(|e| format!("{e:?}"))?;
1205        let mut resp = Map::new();
1206        resp.insert("ok".into(), JsonValue::Bool(true));
1207        resp.insert("collection".into(), JsonValue::String(name.to_string()));
1208        json_to_string(&JsonValue::Object(resp)).map_err(|e| format!("serialization error: {}", e))
1209    }
1210
1211    fn tool_drop_collection(&self, args: &JsonValue) -> Result<String, String> {
1212        let name = get_str_field(args, "name")?;
1213        self.runtime
1214            .db()
1215            .store()
1216            .drop_collection(name)
1217            .map_err(|e| format!("{e:?}"))?;
1218        let mut resp = Map::new();
1219        resp.insert("ok".into(), JsonValue::Bool(true));
1220        resp.insert("dropped".into(), JsonValue::String(name.to_string()));
1221        json_to_string(&JsonValue::Object(resp)).map_err(|e| format!("serialization error: {}", e))
1222    }
1223}
1224
1225// ------------------------------------------------------------------
1226// Helpers
1227// ------------------------------------------------------------------
1228
1229fn format_mcp_ask_parse_error(err: crate::runtime::ai::mcp_ask_tool::ParseError) -> String {
1230    use crate::runtime::ai::mcp_ask_tool::ParseError;
1231
1232    match err {
1233        ParseError::NotAnObject => "arguments must be an object".to_string(),
1234        ParseError::MissingQuestion => "missing required field 'question'".to_string(),
1235        ParseError::QuestionWrongType => "field 'question' must be a string".to_string(),
1236        ParseError::WrongType { path, expected } => {
1237            format!("{path} must be {expected}")
1238        }
1239        ParseError::OutOfRange { path, detail } => {
1240            format!("{path} out of range: {detail}")
1241        }
1242        ParseError::CacheAndNocache => {
1243            "options.cache and options.nocache are mutually exclusive".to_string()
1244        }
1245        ParseError::UnknownOption { path } => format!("unknown option {path}"),
1246    }
1247}
1248
1249fn parse_direction(s: Option<&str>) -> RuntimeGraphDirection {
1250    match s {
1251        Some("incoming") => RuntimeGraphDirection::Incoming,
1252        Some("both") => RuntimeGraphDirection::Both,
1253        _ => RuntimeGraphDirection::Outgoing,
1254    }
1255}
1256
1257/// Parse optional metadata from an `args` JSON object.
1258fn parse_metadata_arg(
1259    args: &JsonValue,
1260) -> Result<Vec<(String, crate::storage::unified::MetadataValue)>, String> {
1261    match args.get("metadata").and_then(|v| v.as_object()) {
1262        Some(obj) => {
1263            let mut out = Vec::with_capacity(obj.len());
1264            for (key, value) in obj {
1265                let mv = crate::application::entity::json_to_metadata_value(value)
1266                    .map_err(|e| format!("{}", e))?;
1267                out.push((key.clone(), mv));
1268            }
1269            Ok(out)
1270        }
1271        None => Ok(vec![]),
1272    }
1273}
1274
1275fn parse_string_array_arg(args: &JsonValue, field: &str) -> Result<Vec<String>, String> {
1276    match args.get(field) {
1277        None | Some(JsonValue::Null) => Ok(Vec::new()),
1278        Some(JsonValue::Array(values)) => values
1279            .iter()
1280            .map(|value| {
1281                value
1282                    .as_str()
1283                    .map(ToOwned::to_owned)
1284                    .ok_or_else(|| format!("field '{field}' must be an array of strings"))
1285            })
1286            .collect(),
1287        _ => Err(format!("field '{field}' must be an array of strings")),
1288    }
1289}
1290
1291fn mcp_keyed_ident(value: &str) -> Result<String, String> {
1292    if !value.is_empty()
1293        && value
1294            .bytes()
1295            .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'.')
1296    {
1297        Ok(value.to_string())
1298    } else {
1299        Err(
1300            "keyed collection and key names must use letters, numbers, underscores, or dots"
1301                .to_string(),
1302        )
1303    }
1304}
1305
1306fn mcp_value_literal(value: &JsonValue) -> Result<String, String> {
1307    match value {
1308        JsonValue::String(value) => Ok(format!("'{}'", value.replace('\'', "''"))),
1309        JsonValue::Number(value) => Ok(value.to_string()),
1310        JsonValue::Bool(value) => Ok(value.to_string()),
1311        JsonValue::Null => Ok("NULL".to_string()),
1312        JsonValue::Array(_) | JsonValue::Object(_) => {
1313            json_to_string(value).map_err(|err| format!("serialization error: {err}"))
1314        }
1315    }
1316}
1317
1318fn append_mcp_tags_clause(sql: &mut String, tags: &[String]) {
1319    if tags.is_empty() {
1320        return;
1321    }
1322    sql.push_str(" TAGS [");
1323    for (index, tag) in tags.iter().enumerate() {
1324        if index > 0 {
1325            sql.push_str(", ");
1326        }
1327        sql.push('\'');
1328        sql.push_str(&tag.replace('\'', "''"));
1329        sql.push('\'');
1330    }
1331    sql.push(']');
1332}
1333
1334fn reject_mcp_volatile_options(args: &JsonValue, domain: &str) -> Result<(), String> {
1335    for field in ["ttl", "ttl_ms", "expire", "expire_ms", "expires_at"] {
1336        if args.get(field).is_some() {
1337            return Err(format!(
1338                "{domain} does not support TTL or expiration options"
1339            ));
1340        }
1341    }
1342    Ok(())
1343}
1344
1345// Convert a storage Value to JSON (local helper to avoid visibility issues).
1346fn get_str_field<'a>(args: &'a JsonValue, field: &str) -> Result<&'a str, String> {
1347    args.get(field)
1348        .and_then(|v| v.as_str())
1349        .ok_or_else(|| format!("missing '{field}'"))
1350}
1351
1352// Auth tool implementations
1353impl McpServer {
1354    fn tool_auth_bootstrap(&self, args: &JsonValue) -> Result<String, String> {
1355        let username = get_str_field(args, "username")?;
1356        let password = get_str_field(args, "password")?;
1357
1358        let br = self
1359            .auth_store
1360            .bootstrap(username, password)
1361            .map_err(|e| e.to_string())?;
1362
1363        let mut result = Map::new();
1364        result.insert("ok".into(), JsonValue::Bool(true));
1365        result.insert("username".into(), JsonValue::String(br.user.username));
1366        result.insert(
1367            "role".into(),
1368            JsonValue::String(br.user.role.as_str().into()),
1369        );
1370        result.insert("api_key".into(), JsonValue::String(br.api_key.key));
1371        result.insert("api_key_name".into(), JsonValue::String(br.api_key.name));
1372        if let Some(cert) = br.certificate {
1373            result.insert("certificate".into(), JsonValue::String(cert));
1374            result.insert(
1375                "message".into(),
1376                JsonValue::String(
1377                    "Save this certificate — it is the ONLY way to unseal the vault after restart."
1378                        .into(),
1379                ),
1380            );
1381        } else {
1382            result.insert(
1383                "message".into(),
1384                JsonValue::String(
1385                    "First admin user created. Save the API key — it won't be shown again.".into(),
1386                ),
1387            );
1388        }
1389        json_to_string(&JsonValue::Object(result))
1390    }
1391
1392    fn tool_auth_create_user(&self, args: &JsonValue) -> Result<String, String> {
1393        let username = get_str_field(args, "username")?;
1394        let password = get_str_field(args, "password")?;
1395        let role_str = get_str_field(args, "role")?;
1396        let role = Role::from_str(role_str).ok_or_else(|| format!("invalid role: {role_str}"))?;
1397
1398        self.auth_store
1399            .create_user(username, password, role)
1400            .map_err(|e| e.to_string())?;
1401
1402        let mut result = Map::new();
1403        result.insert("ok".into(), JsonValue::Bool(true));
1404        result.insert("username".into(), JsonValue::String(username.into()));
1405        result.insert("role".into(), JsonValue::String(role.as_str().into()));
1406        json_to_string(&JsonValue::Object(result))
1407    }
1408
1409    fn tool_auth_login(&self, args: &JsonValue) -> Result<String, String> {
1410        let username = get_str_field(args, "username")?;
1411        let password = get_str_field(args, "password")?;
1412
1413        let session = self
1414            .auth_store
1415            .authenticate(username, password)
1416            .map_err(|e| e.to_string())?;
1417
1418        let mut result = Map::new();
1419        result.insert("ok".into(), JsonValue::Bool(true));
1420        result.insert("token".into(), JsonValue::String(session.token));
1421        result.insert("username".into(), JsonValue::String(session.username));
1422        result.insert(
1423            "role".into(),
1424            JsonValue::String(session.role.as_str().into()),
1425        );
1426        result.insert(
1427            "expires_at".into(),
1428            JsonValue::Number(session.expires_at as f64),
1429        );
1430        json_to_string(&JsonValue::Object(result))
1431    }
1432
1433    fn tool_auth_create_api_key(&self, args: &JsonValue) -> Result<String, String> {
1434        let username = get_str_field(args, "username")?;
1435        let name = get_str_field(args, "name")?;
1436        let role_str = get_str_field(args, "role")?;
1437        let role = Role::from_str(role_str).ok_or_else(|| format!("invalid role: {role_str}"))?;
1438
1439        let key = self
1440            .auth_store
1441            .create_api_key(username, name, role)
1442            .map_err(|e| e.to_string())?;
1443
1444        let mut result = Map::new();
1445        result.insert("ok".into(), JsonValue::Bool(true));
1446        result.insert("key".into(), JsonValue::String(key.key));
1447        result.insert("name".into(), JsonValue::String(key.name));
1448        result.insert("role".into(), JsonValue::String(key.role.as_str().into()));
1449        json_to_string(&JsonValue::Object(result))
1450    }
1451
1452    fn tool_auth_list_users(&self) -> Result<String, String> {
1453        let users = self.auth_store.list_users();
1454        let arr: Vec<JsonValue> = users
1455            .into_iter()
1456            .map(|u| {
1457                let mut obj = Map::new();
1458                obj.insert("username".into(), JsonValue::String(u.username));
1459                obj.insert("role".into(), JsonValue::String(u.role.as_str().into()));
1460                obj.insert("enabled".into(), JsonValue::Bool(u.enabled));
1461                obj.insert(
1462                    "api_key_count".into(),
1463                    JsonValue::Number(u.api_keys.len() as f64),
1464                );
1465                JsonValue::Object(obj)
1466            })
1467            .collect();
1468        json_to_string(&JsonValue::Array(arr))
1469    }
1470}
1471
1472#[cfg(test)]
1473mod tests {
1474    use super::*;
1475    use std::io::{Read, Write};
1476    use std::net::{SocketAddr, TcpListener, TcpStream};
1477    use std::sync::atomic::{AtomicBool, Ordering};
1478    use std::sync::{Arc, Mutex};
1479    use std::thread::{self, JoinHandle};
1480    use std::time::Duration;
1481
1482    static ASK_ENV_LOCK: Mutex<()> = Mutex::new(());
1483
1484    fn make_server() -> McpServer {
1485        let rt = RedDBRuntime::in_memory().expect("in-memory runtime");
1486        McpServer::new(rt)
1487    }
1488
1489    fn parse_json(s: &str) -> JsonValue {
1490        json_from_str(s).expect("valid json")
1491    }
1492
1493    #[test]
1494    fn tools_list_registers_reddb_ask_descriptor() {
1495        let srv = make_server();
1496        let response = srv.handle_tools_list(Some(&JsonValue::Number(1.0)));
1497        let parsed = parse_json(&response);
1498        let tools = parsed
1499            .get("result")
1500            .and_then(|result| result.get("tools"))
1501            .and_then(JsonValue::as_array)
1502            .expect("tools array");
1503
1504        let ask = tools
1505            .iter()
1506            .find(|tool| tool.get("name").and_then(JsonValue::as_str) == Some("reddb.ask"))
1507            .expect("reddb.ask registered");
1508
1509        let desc = ask
1510            .get("description")
1511            .and_then(JsonValue::as_str)
1512            .expect("description");
1513        assert!(desc.contains("citations"), "description: {desc}");
1514        assert!(desc.contains("sources_flat"), "description: {desc}");
1515        assert!(desc.contains("URN"), "description: {desc}");
1516
1517        let options = ask
1518            .get("inputSchema")
1519            .and_then(|schema| schema.get("properties"))
1520            .and_then(|props| props.get("options"))
1521            .and_then(|opts| opts.get("properties"))
1522            .and_then(JsonValue::as_object)
1523            .expect("options properties");
1524        for key in [
1525            "strict",
1526            "using",
1527            "model",
1528            "limit",
1529            "min_score",
1530            "depth",
1531            "temperature",
1532            "seed",
1533            "cache",
1534            "nocache",
1535        ] {
1536            assert!(
1537                options.contains_key(key),
1538                "missing option {key} in {options:?}"
1539            );
1540        }
1541    }
1542
1543    #[test]
1544    fn tools_call_reddb_ask_uses_typed_argument_parser() {
1545        let srv = make_server();
1546        let params = parse_json(
1547            r#"{
1548                "name": "reddb.ask",
1549                "arguments": {
1550                    "question": "what cites this?",
1551                    "options": { "tempurature": 0.2 }
1552                }
1553            }"#,
1554        );
1555
1556        let response = srv.handle_tools_call(Some(&JsonValue::Number(1.0)), Some(&params));
1557        let parsed = parse_json(&response);
1558        let result = parsed.get("result").expect("result");
1559        assert_eq!(
1560            result.get("isError").and_then(JsonValue::as_bool),
1561            Some(true)
1562        );
1563        let text = result
1564            .get("content")
1565            .and_then(JsonValue::as_array)
1566            .and_then(|content| content.first())
1567            .and_then(|item| item.get("text"))
1568            .and_then(JsonValue::as_str)
1569            .expect("error text");
1570        assert!(text.contains("options.tempurature"), "text: {text}");
1571    }
1572
1573    #[test]
1574    fn tools_call_reddb_ask_returns_canonical_citation_envelope() {
1575        let _guard = ASK_ENV_LOCK.lock().expect("env lock");
1576        let stub = AskStub::start();
1577        let _api_base = EnvVarGuard::set(
1578            "REDDB_OLLAMA_API_BASE",
1579            &format!("http://{}/v1", stub.addr()),
1580        );
1581
1582        let srv = make_server();
1583        srv.tool_query(&parse_json(
1584            r#"{"sql":"CREATE TABLE travel (id INTEGER, passport TEXT, notes TEXT)"}"#,
1585        ))
1586        .expect("ddl ok");
1587        srv.tool_query(&parse_json(
1588            r#"{"sql":"INSERT INTO travel (id, passport, notes) VALUES (1, 'PT-002', 'incident FDD-12313 escalated')"}"#,
1589        ))
1590        .expect("insert ok");
1591
1592        let params = parse_json(
1593            r#"{
1594                "name": "reddb.ask",
1595                "arguments": {
1596                    "question": "passport FDD-12313",
1597                    "options": {
1598                        "strict": false,
1599                        "using": "ollama",
1600                        "model": "mock-ask",
1601                        "limit": 1,
1602                        "min_score": 0,
1603                        "depth": 0,
1604                        "temperature": 0,
1605                        "seed": 0,
1606                        "cache": { "ttl": "5m" }
1607                    }
1608                }
1609            }"#,
1610        );
1611
1612        let response = srv.handle_tools_call(Some(&JsonValue::Number(1.0)), Some(&params));
1613        let parsed = parse_json(&response);
1614        let result = parsed.get("result").expect("result");
1615        assert_ne!(
1616            result.get("isError").and_then(JsonValue::as_bool),
1617            Some(true),
1618            "response: {response}"
1619        );
1620        let text = result
1621            .get("content")
1622            .and_then(JsonValue::as_array)
1623            .and_then(|content| content.first())
1624            .and_then(|item| item.get("text"))
1625            .and_then(JsonValue::as_str)
1626            .expect("tool text");
1627        let envelope = parse_json(text);
1628
1629        assert_eq!(
1630            envelope.get("answer").and_then(JsonValue::as_str),
1631            Some("FDD-12313 escalated [^1].")
1632        );
1633        assert_eq!(
1634            envelope.get("provider").and_then(JsonValue::as_str),
1635            Some("ollama")
1636        );
1637        assert_eq!(
1638            envelope.get("model").and_then(JsonValue::as_str),
1639            Some("mock-ask")
1640        );
1641        assert_eq!(
1642            envelope.get("cache_hit").and_then(JsonValue::as_bool),
1643            Some(false)
1644        );
1645        assert!(envelope
1646            .get("sources_flat")
1647            .and_then(JsonValue::as_array)
1648            .is_some());
1649        assert!(envelope
1650            .get("citations")
1651            .and_then(JsonValue::as_array)
1652            .is_some());
1653        assert!(envelope
1654            .get("validation")
1655            .and_then(JsonValue::as_object)
1656            .is_some());
1657        assert!(
1658            envelope.get("rows").is_none(),
1659            "ASK must not be row-wrapped: {text}"
1660        );
1661    }
1662
1663    #[test]
1664    fn tool_query_without_params_keeps_legacy_path() {
1665        let srv = make_server();
1666        let args = parse_json(r#"{"sql":"SELECT 1 AS one"}"#);
1667        let out = srv.tool_query(&args).expect("query ok");
1668        assert!(out.contains("\"one\""), "expected 'one' column in {out}");
1669    }
1670
1671    #[test]
1672    fn tool_query_binds_int_and_text_params() {
1673        let srv = make_server();
1674        srv.tool_query(&parse_json(
1675            r#"{"sql":"CREATE TABLE mcpp (id INTEGER, name TEXT)"}"#,
1676        ))
1677        .expect("ddl ok");
1678        srv.tool_query(&parse_json(
1679            r#"{"sql":"INSERT INTO mcpp (id, name) VALUES (1, 'Alice')"}"#,
1680        ))
1681        .expect("insert 1");
1682        srv.tool_query(&parse_json(
1683            r#"{"sql":"INSERT INTO mcpp (id, name) VALUES (2, 'Bob')"}"#,
1684        ))
1685        .expect("insert 2");
1686
1687        let out = srv
1688            .tool_query(&parse_json(
1689                r#"{"sql":"SELECT * FROM mcpp WHERE id = $1 AND name = $2","params":[1,"Alice"]}"#,
1690            ))
1691            .expect("query with params ok");
1692        assert!(out.contains("Alice"), "expected Alice in {out}");
1693        assert!(!out.contains("Bob"), "Bob must not match: {out}");
1694    }
1695
1696    #[test]
1697    fn tool_query_params_must_be_array() {
1698        let srv = make_server();
1699        let err = srv
1700            .tool_query(&parse_json(
1701                r#"{"sql":"SELECT 1","params":{"not":"array"}}"#,
1702            ))
1703            .expect_err("must reject non-array params");
1704        assert!(err.contains("array"), "got {err}");
1705    }
1706
1707    #[test]
1708    fn tool_query_param_arity_mismatch_surfaces_error() {
1709        let srv = make_server();
1710        srv.tool_query(&parse_json(r#"{"sql":"CREATE TABLE mcpa (id INTEGER)"}"#))
1711            .expect("ddl ok");
1712        let err = srv
1713            .tool_query(&parse_json(
1714                r#"{"sql":"SELECT * FROM mcpa WHERE id = $1","params":[1,2]}"#,
1715            ))
1716            .expect_err("arity mismatch");
1717        assert!(
1718            err.contains("number of parameters") || err.contains("expects"),
1719            "got {err}"
1720        );
1721    }
1722
1723    #[test]
1724    fn tool_query_vector_param_binds_into_search_similar() {
1725        let srv = make_server();
1726        let out = srv.tool_query(&parse_json(
1727            r#"{"sql":"SEARCH SIMILAR $1 COLLECTION mcpv LIMIT 5","params":[[0.1,0.2,0.3]]}"#,
1728        ));
1729        // The collection doesn't exist; we only need to confirm the
1730        // param-bind path runs (i.e. the error reflects runtime semantics,
1731        // not a `$N` placeholder being unresolved).
1732        if let Err(e) = out {
1733            assert!(
1734                !e.contains("placeholder") && !e.contains("Parameter"),
1735                "param did not bind: {e}"
1736            );
1737        }
1738    }
1739
1740    struct EnvVarGuard {
1741        name: &'static str,
1742        previous: Option<String>,
1743    }
1744
1745    impl EnvVarGuard {
1746        fn set(name: &'static str, value: &str) -> Self {
1747            let previous = std::env::var(name).ok();
1748            std::env::set_var(name, value);
1749            Self { name, previous }
1750        }
1751    }
1752
1753    impl Drop for EnvVarGuard {
1754        fn drop(&mut self) {
1755            if let Some(value) = self.previous.take() {
1756                std::env::set_var(self.name, value);
1757            } else {
1758                std::env::remove_var(self.name);
1759            }
1760        }
1761    }
1762
1763    struct AskStub {
1764        addr: SocketAddr,
1765        shutdown: Arc<AtomicBool>,
1766        handle: Option<JoinHandle<()>>,
1767    }
1768
1769    impl AskStub {
1770        fn start() -> Self {
1771            let listener = TcpListener::bind("127.0.0.1:0").expect("stub bind");
1772            listener
1773                .set_nonblocking(true)
1774                .expect("nonblocking listener");
1775            let addr = listener.local_addr().expect("local addr");
1776            let shutdown = Arc::new(AtomicBool::new(false));
1777            let server_shutdown = Arc::clone(&shutdown);
1778            let handle = thread::spawn(move || {
1779                while !server_shutdown.load(Ordering::Relaxed) {
1780                    match listener.accept() {
1781                        Ok((mut stream, _)) => {
1782                            let request = read_stub_request(&mut stream);
1783                            if request.contains("/embeddings") {
1784                                write_json_response(
1785                                    &mut stream,
1786                                    r#"{"model":"mock-embedding","data":[{"index":0,"embedding":[1,0,0]}],"usage":{"prompt_tokens":3,"total_tokens":3}}"#,
1787                                );
1788                            } else {
1789                                write_json_response(
1790                                    &mut stream,
1791                                    r#"{"model":"mock-ask","choices":[{"message":{"role":"assistant","content":"FDD-12313 escalated [^1]."},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":4,"total_tokens":14}}"#,
1792                                );
1793                            }
1794                        }
1795                        Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => {
1796                            thread::sleep(Duration::from_millis(1));
1797                        }
1798                        Err(_) => break,
1799                    }
1800                }
1801            });
1802
1803            Self {
1804                addr,
1805                shutdown,
1806                handle: Some(handle),
1807            }
1808        }
1809
1810        fn addr(&self) -> SocketAddr {
1811            self.addr
1812        }
1813    }
1814
1815    impl Drop for AskStub {
1816        fn drop(&mut self) {
1817            self.shutdown.store(true, Ordering::Relaxed);
1818            let _ = TcpStream::connect(self.addr);
1819            if let Some(handle) = self.handle.take() {
1820                let _ = handle.join();
1821            }
1822        }
1823    }
1824
1825    fn read_stub_request(stream: &mut TcpStream) -> String {
1826        let _ = stream.set_read_timeout(Some(Duration::from_millis(100)));
1827        let mut buffer = [0_u8; 4096];
1828        let count = stream.read(&mut buffer).unwrap_or(0);
1829        String::from_utf8_lossy(&buffer[..count]).into_owned()
1830    }
1831
1832    fn write_json_response(stream: &mut TcpStream, body: &str) {
1833        let response = format!(
1834            "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
1835            body.len()
1836        );
1837        stream
1838            .write_all(response.as_bytes())
1839            .expect("write stub response");
1840    }
1841}
reddb_server/mcp/server.rs

reddb_server/mcp/
server.rs
