1use 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
33pub struct McpServer {
35 runtime: RedDBRuntime,
36 auth_store: Arc<AuthStore>,
37 initialized: bool,
38}
39
40impl McpServer {
41 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 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 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 None
111 }
112 "tools/list" => Some(self.handle_tools_list(id)),
113 "tools/call" => Some(self.handle_tools_call(id, msg.get("params"))),
114 "resources/list" => Some(self.handle_resources_list(id)),
115 "resources/read" => Some(self.handle_resources_read(id, msg.get("params"))),
116 "ping" => {
117 let mut result = Map::new();
118 result.insert("status".to_string(), JsonValue::String("ok".to_string()));
119 Some(protocol::build_result_message(
120 id,
121 JsonValue::Object(result),
122 ))
123 }
124 _ => Some(protocol::build_error_message(
125 id,
126 -32601,
127 &format!("unknown method: {}", method),
128 )),
129 }
130 }
131
132 fn handle_initialize(&mut self, id: Option<&JsonValue>) -> String {
137 self.initialized = true;
138
139 let mut capabilities = Map::new();
140 {
141 let mut tools_cap = Map::new();
142 tools_cap.insert("listChanged".to_string(), JsonValue::Bool(false));
143 capabilities.insert("tools".to_string(), JsonValue::Object(tools_cap));
144 }
145 {
146 let mut resources_cap = Map::new();
149 resources_cap.insert("subscribe".to_string(), JsonValue::Bool(false));
150 resources_cap.insert("listChanged".to_string(), JsonValue::Bool(false));
151 capabilities.insert("resources".to_string(), JsonValue::Object(resources_cap));
152 }
153
154 let mut server_info = Map::new();
155 server_info.insert(
156 "name".to_string(),
157 JsonValue::String("reddb-mcp".to_string()),
158 );
159 server_info.insert(
160 "version".to_string(),
161 JsonValue::String(env!("CARGO_PKG_VERSION").to_string()),
162 );
163
164 let mut result = Map::new();
165 result.insert(
166 "protocolVersion".to_string(),
167 JsonValue::String("2024-11-05".to_string()),
168 );
169 result.insert("capabilities".to_string(), JsonValue::Object(capabilities));
170 result.insert("serverInfo".to_string(), JsonValue::Object(server_info));
171
172 protocol::build_result_message(id, JsonValue::Object(result))
173 }
174
175 fn handle_tools_list(&self, id: Option<&JsonValue>) -> String {
180 let defs = tools::all_tools();
181 let mut tools_json: Vec<JsonValue> = defs
182 .into_iter()
183 .map(|def| {
184 let mut obj = Map::new();
185 obj.insert("name".to_string(), JsonValue::String(def.name.to_string()));
186 obj.insert(
187 "description".to_string(),
188 JsonValue::String(def.description.to_string()),
189 );
190 obj.insert("inputSchema".to_string(), def.input_schema);
191 JsonValue::Object(obj)
192 })
193 .collect();
194 tools_json.push(crate::runtime::ai::mcp_ask_tool::descriptor());
195
196 let mut result = Map::new();
197 result.insert("tools".to_string(), JsonValue::Array(tools_json));
198 protocol::build_result_message(id, JsonValue::Object(result))
199 }
200
201 fn handle_resources_list(&self, id: Option<&JsonValue>) -> String {
206 let resources: Vec<JsonValue> = tools::knowledge_resources()
207 .iter()
208 .map(|res| {
209 let mut obj = Map::new();
210 obj.insert("uri".to_string(), JsonValue::String(res.uri.to_string()));
211 obj.insert("name".to_string(), JsonValue::String(res.title.to_string()));
212 obj.insert(
213 "description".to_string(),
214 JsonValue::String(res.description.to_string()),
215 );
216 obj.insert(
217 "mimeType".to_string(),
218 JsonValue::String(res.mime_type.to_string()),
219 );
220 JsonValue::Object(obj)
221 })
222 .collect();
223
224 let mut result = Map::new();
225 result.insert("resources".to_string(), JsonValue::Array(resources));
226 protocol::build_result_message(id, JsonValue::Object(result))
227 }
228
229 fn handle_resources_read(&self, id: Option<&JsonValue>, params: Option<&JsonValue>) -> String {
230 let uri = params.and_then(|p| p.get("uri")).and_then(|v| v.as_str());
231 let uri = match uri {
232 Some(u) => u,
233 None => return protocol::build_error_message(id, -32602, "missing resource uri"),
234 };
235
236 let resources = tools::knowledge_resources();
237 let resource = match resources.iter().find(|res| res.uri == uri) {
238 Some(res) => res,
239 None => {
240 return protocol::build_error_message(
241 id,
242 -32602,
243 &format!("unknown resource: {uri}"),
244 );
245 }
246 };
247
248 let mut contents = Map::new();
249 contents.insert(
250 "uri".to_string(),
251 JsonValue::String(resource.uri.to_string()),
252 );
253 contents.insert(
254 "mimeType".to_string(),
255 JsonValue::String(resource.mime_type.to_string()),
256 );
257 contents.insert("text".to_string(), JsonValue::String((resource.body)()));
258
259 let mut result = Map::new();
260 result.insert(
261 "contents".to_string(),
262 JsonValue::Array(vec![JsonValue::Object(contents)]),
263 );
264 protocol::build_result_message(id, JsonValue::Object(result))
265 }
266
267 fn handle_tools_call(&self, id: Option<&JsonValue>, params: Option<&JsonValue>) -> String {
272 let name = params.and_then(|p| p.get("name")).and_then(|v| v.as_str());
273 let name = match name {
274 Some(n) => n,
275 None => {
276 return protocol::build_error_message(id, -32602, "missing tool name");
277 }
278 };
279
280 let empty = JsonValue::Object(Map::new());
281 let args = params.and_then(|p| p.get("arguments")).unwrap_or(&empty);
282
283 let result = match name {
284 "reddb_query" => self.tool_query(args),
285 "reddb_collections" => self.tool_collections(),
286 "reddb_insert_row" => self.tool_insert_row(args),
287 "reddb_insert_node" => self.tool_insert_node(args),
288 "reddb_insert_edge" => self.tool_insert_edge(args),
289 "reddb_insert_vector" => self.tool_insert_vector(args),
290 "reddb_insert_document" => self.tool_insert_document(args),
291 "reddb_kv_get" => self.tool_kv_get(args),
292 "reddb_kv_set" => self.tool_kv_set(args),
293 "reddb_kv_invalidate_tags" => self.tool_kv_invalidate_tags(args),
294 "reddb_config_get" => self.tool_config_get(args),
295 "reddb_config_put" => self.tool_config_put(args),
296 "reddb_config_resolve" => self.tool_config_resolve(args),
297 "reddb_vault_get" => self.tool_vault_get(args),
298 "reddb_vault_put" => self.tool_vault_put(args),
299 "reddb_vault_unseal" => self.tool_vault_unseal(args),
300 "reddb_delete" => self.tool_delete(args),
301 "reddb_search_vector" => self.tool_search_vector(args),
302 "reddb_search_text" => self.tool_search_text(args),
303 "reddb_health" => self.tool_health(),
304 "reddb_graph_traverse" => self.tool_graph_traverse(args),
305 "reddb_graph_shortest_path" => self.tool_graph_shortest_path(args),
306 "reddb_update" => self.tool_update(args),
307 "reddb_scan" => self.tool_scan(args),
308 "reddb_graph_centrality" => self.tool_graph_centrality(args),
309 "reddb_graph_community" => self.tool_graph_community(args),
310 "reddb_graph_components" => self.tool_graph_components(args),
311 "reddb_graph_cycles" => self.tool_graph_cycles(args),
312 "reddb_graph_clustering" => self.tool_graph_clustering(args),
313 "reddb_create_collection" => self.tool_create_collection(args),
314 "reddb_drop_collection" => self.tool_drop_collection(args),
315 "reddb_rql_validate" => self.tool_rql_validate(args),
316 "reddb_rql_explain" => self.tool_rql_explain(args),
317 "reddb_type_of" => self.tool_type_of(args),
318 "reddb_explain_connection" => self.tool_explain_connection(args),
319 "reddb_auth_bootstrap" => self.tool_auth_bootstrap(args),
320 "reddb_auth_create_user" => self.tool_auth_create_user(args),
321 "reddb_auth_login" => self.tool_auth_login(args),
322 "reddb_auth_create_api_key" => self.tool_auth_create_api_key(args),
323 "reddb_auth_list_users" => self.tool_auth_list_users(),
324 crate::runtime::ai::mcp_ask_tool::TOOL_NAME => self.tool_ask(args),
325 _ => Err(format!("unknown tool: {name}")),
326 };
327
328 match result {
329 Ok(text) => {
330 let mut content = Map::new();
331 content.insert("type".to_string(), JsonValue::String("text".to_string()));
332 content.insert("text".to_string(), JsonValue::String(text));
333
334 let mut result_obj = Map::new();
335 result_obj.insert(
336 "content".to_string(),
337 JsonValue::Array(vec![JsonValue::Object(content)]),
338 );
339 protocol::build_result_message(id, JsonValue::Object(result_obj))
340 }
341 Err(err) => {
342 let mut content = Map::new();
343 content.insert("type".to_string(), JsonValue::String("text".to_string()));
344 content.insert("text".to_string(), JsonValue::String(err.clone()));
345
346 let mut result_obj = Map::new();
347 result_obj.insert(
348 "content".to_string(),
349 JsonValue::Array(vec![JsonValue::Object(content)]),
350 );
351 result_obj.insert("isError".to_string(), JsonValue::Bool(true));
352 protocol::build_result_message(id, JsonValue::Object(result_obj))
353 }
354 }
355 }
356
357 fn tool_query(&self, args: &JsonValue) -> Result<String, String> {
362 let sql = args
363 .get("sql")
364 .and_then(|v| v.as_str())
365 .ok_or("missing required field 'sql'")?;
366
367 if let Some(raw_params) = args.get("params") {
371 let arr = raw_params
372 .as_array()
373 .ok_or_else(|| "'params' must be an array".to_string())?;
374 let binds: Vec<Value> = arr
375 .iter()
376 .map(crate::rpc_stdio::json_value_to_schema_value)
377 .collect();
378
379 use crate::storage::query::modes::parse_multi;
380 use crate::storage::query::user_params;
381 let parsed = parse_multi(sql).map_err(|e| format!("{}", e))?;
382 let bound = user_params::bind(&parsed, &binds).map_err(|e| format!("{}", e))?;
383 let result = self
384 .runtime
385 .execute_query_expr(bound)
386 .map_err(|e| format!("{}", e))?;
387 let json = runtime_query_json(&result, &None, &None);
388 return json_to_string(&json).map_err(|e| format!("serialization error: {}", e));
389 }
390
391 let uc = QueryUseCases::new(&self.runtime);
392 let result = uc
393 .execute(ExecuteQueryInput {
394 query: sql.to_string(),
395 })
396 .map_err(|e| format!("{}", e))?;
397
398 let json = runtime_query_json(&result, &None, &None);
399 json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
400 }
401
402 fn tool_ask(&self, args: &JsonValue) -> Result<String, String> {
403 let invocation =
404 crate::runtime::ai::mcp_ask_tool::parse(args).map_err(format_mcp_ask_parse_error)?;
405 let ask = crate::storage::query::ast::AskQuery {
406 explain: false,
407 question: invocation.question,
408 question_param: None,
409 provider: invocation.using,
410 model: invocation.model,
411 depth: invocation.depth.map(|v| v as usize),
412 limit: invocation.limit.map(|v| v as usize),
413 min_score: invocation.min_score.map(|v| v as f32),
414 collection: None,
415 temperature: invocation.temperature.map(|v| v as f32),
416 seed: invocation.seed,
417 strict: invocation.strict.unwrap_or(true),
418 stream: false,
419 cache: if matches!(invocation.nocache, Some(true)) {
420 crate::storage::query::ast::AskCacheClause::NoCache
421 } else if let Some(ttl) = invocation.cache_ttl {
422 crate::storage::query::ast::AskCacheClause::CacheTtl(ttl)
423 } else {
424 crate::storage::query::ast::AskCacheClause::Default
425 },
426 plan_only: false,
427 steps: None,
428 };
429 let result = self
430 .runtime
431 .execute_ask("ASK <mcp>", &ask)
432 .map_err(|e| format!("{}", e))?;
433 let json = crate::rpc_stdio::query_result_to_json(&result);
434 json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
435 }
436
437 fn tool_collections(&self) -> Result<String, String> {
438 let uc = CatalogUseCases::new(&self.runtime);
439 let collections = uc.collections();
440 let json = JsonValue::Array(collections.into_iter().map(JsonValue::String).collect());
441 json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
442 }
443
444 fn tool_insert_row(&self, args: &JsonValue) -> Result<String, String> {
445 let collection = args
446 .get("collection")
447 .and_then(|v| v.as_str())
448 .ok_or("missing required field 'collection'")?;
449 let data = args
450 .get("data")
451 .and_then(|v| v.as_object())
452 .ok_or("missing required field 'data' (must be an object)")?;
453
454 let mut fields = Vec::new();
455 for (key, value) in data {
456 let sv = crate::application::entity::json_to_storage_value(value)
457 .map_err(|e| format!("{}", e))?;
458 fields.push((key.clone(), sv));
459 }
460
461 let metadata = parse_metadata_arg(args)?;
462
463 let uc = EntityUseCases::new(&self.runtime);
464 let output = uc
465 .create_row(CreateRowInput {
466 collection: collection.to_string(),
467 fields,
468 metadata,
469 node_links: vec![],
470 vector_links: vec![],
471 })
472 .map_err(|e| format!("{}", e))?;
473
474 let json = created_entity_output_json(&output);
475 json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
476 }
477
478 fn tool_insert_node(&self, args: &JsonValue) -> Result<String, String> {
479 let collection = args
480 .get("collection")
481 .and_then(|v| v.as_str())
482 .ok_or("missing required field 'collection'")?;
483 let label = args
484 .get("label")
485 .and_then(|v| v.as_str())
486 .ok_or("missing required field 'label'")?;
487 let node_type = args
488 .get("node_type")
489 .and_then(|v| v.as_str())
490 .map(String::from);
491
492 let mut properties = Vec::new();
493 if let Some(props) = args.get("properties").and_then(|v| v.as_object()) {
494 for (key, value) in props {
495 let sv = crate::application::entity::json_to_storage_value(value)
496 .map_err(|e| format!("{}", e))?;
497 properties.push((key.clone(), sv));
498 }
499 }
500
501 let metadata = parse_metadata_arg(args)?;
502
503 let uc = EntityUseCases::new(&self.runtime);
504 let output = uc
505 .create_node(CreateNodeInput {
506 collection: collection.to_string(),
507 label: label.to_string(),
508 node_type,
509 properties,
510 metadata,
511 embeddings: vec![],
512 table_links: vec![],
513 node_links: vec![],
514 })
515 .map_err(|e| format!("{}", e))?;
516
517 let json = created_entity_output_json(&output);
518 json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
519 }
520
521 fn tool_insert_edge(&self, args: &JsonValue) -> Result<String, String> {
522 let collection = args
523 .get("collection")
524 .and_then(|v| v.as_str())
525 .ok_or("missing required field 'collection'")?;
526 let label = args
527 .get("label")
528 .and_then(|v| v.as_str())
529 .ok_or("missing required field 'label'")?;
530 let from_id = args
531 .get("from")
532 .and_then(|v| v.as_u64())
533 .ok_or("missing required field 'from' (integer)")?;
534 let to_id = args
535 .get("to")
536 .and_then(|v| v.as_u64())
537 .ok_or("missing required field 'to' (integer)")?;
538 let weight = args
539 .get("weight")
540 .and_then(|v| v.as_f64())
541 .map(|w| w as f32);
542
543 let mut properties = Vec::new();
544 if let Some(props) = args.get("properties").and_then(|v| v.as_object()) {
545 for (key, value) in props {
546 let sv = crate::application::entity::json_to_storage_value(value)
547 .map_err(|e| format!("{}", e))?;
548 properties.push((key.clone(), sv));
549 }
550 }
551
552 let metadata = parse_metadata_arg(args)?;
553
554 let uc = EntityUseCases::new(&self.runtime);
555 let output = uc
556 .create_edge(CreateEdgeInput {
557 collection: collection.to_string(),
558 label: label.to_string(),
559 from: EntityId::new(from_id),
560 to: EntityId::new(to_id),
561 weight,
562 properties,
563 metadata,
564 })
565 .map_err(|e| format!("{}", e))?;
566
567 let json = created_entity_output_json(&output);
568 json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
569 }
570
571 fn tool_insert_vector(&self, args: &JsonValue) -> Result<String, String> {
572 let collection = args
573 .get("collection")
574 .and_then(|v| v.as_str())
575 .ok_or("missing required field 'collection'")?;
576 let dense_arr = args
577 .get("dense")
578 .and_then(|v| v.as_array())
579 .ok_or("missing required field 'dense' (array of numbers)")?;
580
581 let mut dense = Vec::with_capacity(dense_arr.len());
582 for v in dense_arr {
583 dense.push(
584 v.as_f64()
585 .ok_or("'dense' array must contain only numbers")? as f32,
586 );
587 }
588 if dense.is_empty() {
589 return Err("'dense' vector cannot be empty".to_string());
590 }
591
592 let content = args
593 .get("content")
594 .and_then(|v| v.as_str())
595 .map(String::from);
596 let metadata = parse_metadata_arg(args)?;
597
598 let uc = EntityUseCases::new(&self.runtime);
599 let output = uc
600 .create_vector(CreateVectorInput {
601 collection: collection.to_string(),
602 dense,
603 content,
604 metadata,
605 link_row: None,
606 link_node: None,
607 })
608 .map_err(|e| format!("{}", e))?;
609
610 let json = created_entity_output_json(&output);
611 json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
612 }
613
614 fn tool_insert_document(&self, args: &JsonValue) -> Result<String, String> {
615 let collection = args
616 .get("collection")
617 .and_then(|v| v.as_str())
618 .ok_or("missing required field 'collection'")?;
619 let body = args.get("body").ok_or("missing required field 'body'")?;
620
621 let metadata = parse_metadata_arg(args)?;
622
623 let uc = EntityUseCases::new(&self.runtime);
624 let output = uc
625 .create_document(CreateDocumentInput {
626 collection: collection.to_string(),
627 body: body.clone(),
628 metadata,
629 node_links: vec![],
630 vector_links: vec![],
631 })
632 .map_err(|e| format!("{}", e))?;
633
634 let json = created_entity_output_json(&output);
635 json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
636 }
637
638 fn tool_kv_get(&self, args: &JsonValue) -> Result<String, String> {
639 let collection = args
640 .get("collection")
641 .and_then(|v| v.as_str())
642 .ok_or("missing required field 'collection'")?;
643 let key = args
644 .get("key")
645 .and_then(|v| v.as_str())
646 .ok_or("missing required field 'key'")?;
647
648 let uc = EntityUseCases::new(&self.runtime);
649 match uc.get_kv(collection, key).map_err(|e| format!("{}", e))? {
650 Some((value, entity_id)) => {
651 let mut obj = Map::new();
652 obj.insert("found".to_string(), JsonValue::Bool(true));
653 obj.insert("key".to_string(), JsonValue::String(key.to_string()));
654 obj.insert("value".to_string(), storage_value_to_json(&value));
655 obj.insert("rid".to_string(), JsonValue::Number(entity_id.raw() as f64));
656 obj.insert("kind".to_string(), JsonValue::String("kv".to_string()));
657 json_to_string(&JsonValue::Object(obj))
658 .map_err(|e| format!("serialization error: {}", e))
659 }
660 None => {
661 let mut obj = Map::new();
662 obj.insert("found".to_string(), JsonValue::Bool(false));
663 obj.insert("key".to_string(), JsonValue::String(key.to_string()));
664 json_to_string(&JsonValue::Object(obj))
665 .map_err(|e| format!("serialization error: {}", e))
666 }
667 }
668 }
669
670 fn tool_kv_set(&self, args: &JsonValue) -> Result<String, String> {
671 let collection = args
672 .get("collection")
673 .and_then(|v| v.as_str())
674 .ok_or("missing required field 'collection'")?;
675 let key = args
676 .get("key")
677 .and_then(|v| v.as_str())
678 .ok_or("missing required field 'key'")?;
679 let value_arg = args.get("value").ok_or("missing required field 'value'")?;
680
681 let sv = crate::application::entity::json_to_storage_value(value_arg)
682 .map_err(|e| format!("{}", e))?;
683
684 let metadata = parse_metadata_arg(args)?;
685
686 let tags = parse_string_array_arg(args, "tags")?;
687 let ops = crate::runtime::impl_kv::KvAtomicOps::new(&self.runtime);
688 let (_, id) = ops
689 .set_with_tags_and_metadata(collection, key, sv, None, &tags, false, metadata)
690 .map_err(|e| format!("{}", e))?;
691
692 let mut obj = Map::new();
693 obj.insert("ok".to_string(), JsonValue::Bool(true));
694 obj.insert("rid".to_string(), JsonValue::Number(id.raw() as f64));
695 obj.insert("kind".to_string(), JsonValue::String("kv".to_string()));
696 obj.insert(
697 "tags".to_string(),
698 JsonValue::Array(tags.into_iter().map(JsonValue::String).collect()),
699 );
700 json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
701 }
702
703 fn tool_kv_invalidate_tags(&self, args: &JsonValue) -> Result<String, String> {
704 let collection = args
705 .get("collection")
706 .and_then(|v| v.as_str())
707 .ok_or("missing required field 'collection'")?;
708 let tags = parse_string_array_arg(args, "tags")?;
709 if tags.is_empty() {
710 return Err("missing required field 'tags'".to_string());
711 }
712 let ops = crate::runtime::impl_kv::KvAtomicOps::new(&self.runtime);
713 let count = ops
714 .invalidate_tags(collection, &tags)
715 .map_err(|e| format!("{}", e))?;
716
717 let mut obj = Map::new();
718 obj.insert("ok".to_string(), JsonValue::Bool(true));
719 obj.insert("invalidated".to_string(), JsonValue::Number(count as f64));
720 obj.insert(
721 "tags".to_string(),
722 JsonValue::Array(tags.into_iter().map(JsonValue::String).collect()),
723 );
724 json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
725 }
726
727 fn tool_config_get(&self, args: &JsonValue) -> Result<String, String> {
728 let collection = mcp_keyed_ident(get_str_field(args, "collection")?)?;
729 let key = mcp_keyed_ident(get_str_field(args, "key")?)?;
730 self.tool_keyed_query(format!("GET CONFIG {collection} {key}"))
731 }
732
733 fn tool_config_put(&self, args: &JsonValue) -> Result<String, String> {
734 reject_mcp_volatile_options(args, "CONFIG")?;
735 let collection = mcp_keyed_ident(get_str_field(args, "collection")?)?;
736 let key = mcp_keyed_ident(get_str_field(args, "key")?)?;
737 let tags = parse_string_array_arg(args, "tags")?;
738 let literal = if let Some(secret_ref) = args.get("secret_ref") {
739 let object = secret_ref
740 .as_object()
741 .ok_or("field 'secret_ref' must be an object")?;
742 let ref_collection = object
743 .get("collection")
744 .and_then(|v| v.as_str())
745 .ok_or_else(|| "secret_ref.collection is required".to_string())
746 .and_then(mcp_keyed_ident)?;
747 let ref_key = object
748 .get("key")
749 .and_then(|v| v.as_str())
750 .ok_or_else(|| "secret_ref.key is required".to_string())
751 .and_then(mcp_keyed_ident)?;
752 format!("SECRET_REF(vault, {ref_collection}.{ref_key})")
753 } else {
754 mcp_value_literal(args.get("value").ok_or("missing required field 'value'")?)?
755 };
756 let mut sql = format!("PUT CONFIG {collection} {key} = {literal}");
757 append_mcp_tags_clause(&mut sql, &tags);
758 self.tool_keyed_query(sql)
759 }
760
761 fn tool_config_resolve(&self, args: &JsonValue) -> Result<String, String> {
762 let collection = mcp_keyed_ident(get_str_field(args, "collection")?)?;
763 let key = mcp_keyed_ident(get_str_field(args, "key")?)?;
764 self.tool_keyed_query(format!("RESOLVE CONFIG {collection} {key}"))
765 }
766
767 fn tool_vault_get(&self, args: &JsonValue) -> Result<String, String> {
768 let collection = mcp_keyed_ident(get_str_field(args, "collection")?)?;
769 let key = mcp_keyed_ident(get_str_field(args, "key")?)?;
770 self.tool_keyed_query(format!("VAULT GET {collection}.{key}"))
771 }
772
773 fn tool_vault_put(&self, args: &JsonValue) -> Result<String, String> {
774 reject_mcp_volatile_options(args, "VAULT")?;
775 let collection = mcp_keyed_ident(get_str_field(args, "collection")?)?;
776 let key = mcp_keyed_ident(get_str_field(args, "key")?)?;
777 let value = mcp_value_literal(args.get("value").ok_or("missing required field 'value'")?)?;
778 let tags = parse_string_array_arg(args, "tags")?;
779 let mut sql = format!("VAULT PUT {collection}.{key} = {value}");
780 append_mcp_tags_clause(&mut sql, &tags);
781 self.tool_keyed_query(sql)
782 }
783
784 fn tool_vault_unseal(&self, args: &JsonValue) -> Result<String, String> {
785 let collection = mcp_keyed_ident(get_str_field(args, "collection")?)?;
786 let key = mcp_keyed_ident(get_str_field(args, "key")?)?;
787 self.tool_keyed_query(format!("UNSEAL VAULT {collection}.{key}"))
788 }
789
790 fn tool_keyed_query(&self, sql: String) -> Result<String, String> {
791 let uc = QueryUseCases::new(&self.runtime);
792 let result = uc
793 .execute(ExecuteQueryInput { query: sql })
794 .map_err(|e| format!("{}", e))?;
795 let json = runtime_query_json(&result, &None, &None);
796 json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
797 }
798
799 fn tool_delete(&self, args: &JsonValue) -> Result<String, String> {
800 let collection = args
801 .get("collection")
802 .and_then(|v| v.as_str())
803 .ok_or("missing required field 'collection'")?;
804 let id = args
805 .get("id")
806 .and_then(|v| v.as_u64())
807 .ok_or("missing required field 'id' (integer)")?;
808
809 let uc = EntityUseCases::new(&self.runtime);
810 let output = uc
811 .delete(DeleteEntityInput {
812 collection: collection.to_string(),
813 id: EntityId::new(id),
814 })
815 .map_err(|e| format!("{}", e))?;
816
817 let mut obj = Map::new();
818 obj.insert("deleted".to_string(), JsonValue::Bool(output.deleted));
819 obj.insert("id".to_string(), JsonValue::Number(output.id.raw() as f64));
820 json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
821 }
822
823 fn tool_search_vector(&self, args: &JsonValue) -> Result<String, String> {
824 let collection = args
825 .get("collection")
826 .and_then(|v| v.as_str())
827 .ok_or("missing required field 'collection'")?;
828 let vector_arr = args
829 .get("vector")
830 .and_then(|v| v.as_array())
831 .ok_or("missing required field 'vector' (array of numbers)")?;
832
833 let mut vector = Vec::with_capacity(vector_arr.len());
834 for v in vector_arr {
835 vector.push(
836 v.as_f64()
837 .ok_or("'vector' array must contain only numbers")? as f32,
838 );
839 }
840 let k = args
841 .get("k")
842 .and_then(|v| v.as_u64())
843 .map(|v| v as usize)
844 .unwrap_or(10);
845 let min_score = args
846 .get("min_score")
847 .and_then(|v| v.as_f64())
848 .map(|v| v as f32)
849 .unwrap_or(0.0);
850
851 let uc = QueryUseCases::new(&self.runtime);
852 let results = uc
853 .search_similar(SearchSimilarInput {
854 collection: collection.to_string(),
855 vector,
856 k,
857 min_score,
858 text: None,
859 provider: None,
860 })
861 .map_err(|e| format!("{}", e))?;
862
863 let items: Vec<JsonValue> = results
864 .iter()
865 .map(|r| {
866 let mut obj = Map::new();
867 obj.insert(
868 "rid".to_string(),
869 JsonValue::Number(r.entity_id.raw() as f64),
870 );
871 obj.insert("kind".to_string(), JsonValue::String("vector".to_string()));
872 obj.insert("score".to_string(), JsonValue::Number(r.score as f64));
873 obj.insert("distance".to_string(), JsonValue::Number(r.distance as f64));
874 JsonValue::Object(obj)
875 })
876 .collect();
877
878 let mut obj = Map::new();
879 obj.insert("count".to_string(), JsonValue::Number(items.len() as f64));
880 obj.insert("results".to_string(), JsonValue::Array(items));
881 json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
882 }
883
884 fn tool_search_text(&self, args: &JsonValue) -> Result<String, String> {
885 let query = args
886 .get("query")
887 .and_then(|v| v.as_str())
888 .ok_or("missing required field 'query'")?;
889
890 let collections = args
891 .get("collections")
892 .and_then(|v| v.as_array())
893 .map(|arr| {
894 arr.iter()
895 .filter_map(|v| v.as_str().map(String::from))
896 .collect::<Vec<_>>()
897 });
898 let limit = args
899 .get("limit")
900 .and_then(|v| v.as_u64())
901 .map(|v| v as usize);
902 let fuzzy = args.get("fuzzy").and_then(|v| v.as_bool()).unwrap_or(false);
903
904 let uc = QueryUseCases::new(&self.runtime);
905 let result = uc
906 .search_text(SearchTextInput {
907 query: query.to_string(),
908 collections,
909 entity_types: None,
910 capabilities: None,
911 fields: None,
912 limit,
913 fuzzy,
914 })
915 .map_err(|e| format!("{}", e))?;
916
917 let items: Vec<JsonValue> = result
918 .matches
919 .iter()
920 .map(|m| {
921 let mut obj = Map::new();
922 obj.insert(
923 "rid".to_string(),
924 JsonValue::Number(m.entity.id.raw() as f64),
925 );
926 obj.insert(
927 "kind".to_string(),
928 JsonValue::String(m.entity.kind.storage_type().to_string()),
929 );
930 obj.insert("score".to_string(), JsonValue::Number(m.score as f64));
931 JsonValue::Object(obj)
932 })
933 .collect();
934
935 let mut obj = Map::new();
936 obj.insert("count".to_string(), JsonValue::Number(items.len() as f64));
937 obj.insert("results".to_string(), JsonValue::Array(items));
938 json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
939 }
940
941 fn tool_health(&self) -> Result<String, String> {
942 let uc = CatalogUseCases::new(&self.runtime);
943 let stats = uc.stats();
944 let json = runtime_stats_json(&stats);
945 json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
946 }
947
948 fn tool_graph_traverse(&self, args: &JsonValue) -> Result<String, String> {
949 let source = args
950 .get("source")
951 .and_then(|v| v.as_str())
952 .ok_or("missing required field 'source'")?;
953 let direction = parse_direction(args.get("direction").and_then(|v| v.as_str()));
954 let max_depth = args
955 .get("max_depth")
956 .and_then(|v| v.as_u64())
957 .map(|v| v as usize)
958 .unwrap_or(3);
959 let strategy = match args.get("strategy").and_then(|v| v.as_str()) {
960 Some("dfs") => RuntimeGraphTraversalStrategy::Dfs,
961 _ => RuntimeGraphTraversalStrategy::Bfs,
962 };
963
964 let uc = GraphUseCases::new(&self.runtime);
965 let result = uc
966 .traverse(GraphTraversalInput {
967 source: source.to_string(),
968 direction,
969 max_depth,
970 strategy,
971 edge_labels: None,
972 projection: None,
973 })
974 .map_err(|e| format!("{}", e))?;
975
976 let visits: Vec<JsonValue> = result
977 .visits
978 .iter()
979 .map(|v| {
980 let mut obj = Map::new();
981 obj.insert("depth".to_string(), JsonValue::Number(v.depth as f64));
982 obj.insert("node_id".to_string(), JsonValue::String(v.node.id.clone()));
983 obj.insert("label".to_string(), JsonValue::String(v.node.label.clone()));
984 obj.insert(
985 "node_type".to_string(),
986 JsonValue::String(v.node.node_type.clone()),
987 );
988 JsonValue::Object(obj)
989 })
990 .collect();
991
992 let edges: Vec<JsonValue> = result
993 .edges
994 .iter()
995 .map(|e| {
996 let mut obj = Map::new();
997 obj.insert("source".to_string(), JsonValue::String(e.source.clone()));
998 obj.insert("target".to_string(), JsonValue::String(e.target.clone()));
999 obj.insert(
1000 "edge_type".to_string(),
1001 JsonValue::String(e.edge_type.clone()),
1002 );
1003 obj.insert("weight".to_string(), JsonValue::Number(e.weight as f64));
1004 JsonValue::Object(obj)
1005 })
1006 .collect();
1007
1008 let mut obj = Map::new();
1009 obj.insert(
1010 "source".to_string(),
1011 JsonValue::String(result.source.clone()),
1012 );
1013 obj.insert(
1014 "visit_count".to_string(),
1015 JsonValue::Number(visits.len() as f64),
1016 );
1017 obj.insert("visits".to_string(), JsonValue::Array(visits));
1018 obj.insert("edges".to_string(), JsonValue::Array(edges));
1019 json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
1020 }
1021
1022 fn tool_graph_shortest_path(&self, args: &JsonValue) -> Result<String, String> {
1023 let source = args
1024 .get("source")
1025 .and_then(|v| v.as_str())
1026 .ok_or("missing required field 'source'")?;
1027 let target = args
1028 .get("target")
1029 .and_then(|v| v.as_str())
1030 .ok_or("missing required field 'target'")?;
1031 let direction = parse_direction(args.get("direction").and_then(|v| v.as_str()));
1032 let algorithm = match args.get("algorithm").and_then(|v| v.as_str()) {
1033 Some("astar") | Some("a*") => RuntimeGraphPathAlgorithm::AStar,
1034 Some("bellman_ford") | Some("bellmanford") => RuntimeGraphPathAlgorithm::BellmanFord,
1035 Some("dijkstra") => RuntimeGraphPathAlgorithm::Dijkstra,
1036 _ => RuntimeGraphPathAlgorithm::Bfs,
1037 };
1038
1039 let uc = GraphUseCases::new(&self.runtime);
1040 let result = uc
1041 .shortest_path(GraphShortestPathInput {
1042 source: source.to_string(),
1043 target: target.to_string(),
1044 direction,
1045 algorithm,
1046 edge_labels: None,
1047 projection: None,
1048 })
1049 .map_err(|e| format!("{}", e))?;
1050
1051 let mut obj = Map::new();
1052 obj.insert(
1053 "source".to_string(),
1054 JsonValue::String(result.source.clone()),
1055 );
1056 obj.insert(
1057 "target".to_string(),
1058 JsonValue::String(result.target.clone()),
1059 );
1060 obj.insert(
1061 "nodes_visited".to_string(),
1062 JsonValue::Number(result.nodes_visited as f64),
1063 );
1064
1065 match &result.path {
1066 Some(path) => {
1067 obj.insert("found".to_string(), JsonValue::Bool(true));
1068 obj.insert(
1069 "hop_count".to_string(),
1070 JsonValue::Number(path.hop_count as f64),
1071 );
1072 obj.insert(
1073 "total_weight".to_string(),
1074 JsonValue::Number(path.total_weight),
1075 );
1076 let nodes_json: Vec<JsonValue> = path
1077 .nodes
1078 .iter()
1079 .map(|n| {
1080 let mut nobj = Map::new();
1081 nobj.insert("id".to_string(), JsonValue::String(n.id.clone()));
1082 nobj.insert("label".to_string(), JsonValue::String(n.label.clone()));
1083 JsonValue::Object(nobj)
1084 })
1085 .collect();
1086 obj.insert("nodes".to_string(), JsonValue::Array(nodes_json));
1087 }
1088 None => {
1089 obj.insert("found".to_string(), JsonValue::Bool(false));
1090 }
1091 }
1092
1093 json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
1094 }
1095
1096 fn tool_update(&self, args: &JsonValue) -> Result<String, String> {
1097 let collection = get_str_field(args, "collection")?;
1098 let set_obj = args.get("set").ok_or("missing 'set'")?;
1099 let where_clause = args
1100 .get("where_filter")
1101 .and_then(|v| v.as_str())
1102 .unwrap_or("");
1103
1104 let mut sql = format!("UPDATE {} SET ", collection);
1106 if let Some(obj) = set_obj.as_object() {
1107 let assignments: Vec<String> = obj
1108 .iter()
1109 .map(|(k, v)| {
1110 let val_str = match v {
1111 JsonValue::String(s) => format!("'{}'", s),
1112 JsonValue::Number(n) => n.to_string(),
1113 JsonValue::Bool(b) => b.to_string(),
1114 _ => format!("'{}'", v),
1115 };
1116 format!("{} = {}", k, val_str)
1117 })
1118 .collect();
1119 sql.push_str(&assignments.join(", "));
1120 } else {
1121 return Err("'set' must be a JSON object".to_string());
1122 }
1123 if !where_clause.is_empty() {
1124 sql.push_str(&format!(" WHERE {}", where_clause));
1125 }
1126
1127 let uc = QueryUseCases::new(&self.runtime);
1128 let result = uc
1129 .execute(ExecuteQueryInput { query: sql })
1130 .map_err(|e| format!("{}", e))?;
1131
1132 let mut resp = Map::new();
1133 resp.insert("ok".into(), JsonValue::Bool(true));
1134 resp.insert(
1135 "affected_rows".into(),
1136 JsonValue::Number(result.affected_rows as f64),
1137 );
1138 json_to_string(&JsonValue::Object(resp)).map_err(|e| format!("serialization error: {}", e))
1139 }
1140
1141 fn tool_scan(&self, args: &JsonValue) -> Result<String, String> {
1142 let collection = get_str_field(args, "collection")?;
1143 let limit = args
1144 .get("limit")
1145 .and_then(|v| v.as_u64())
1146 .map(|v| v as usize)
1147 .unwrap_or(10);
1148 let offset = args
1149 .get("offset")
1150 .and_then(|v| v.as_u64())
1151 .map(|v| v as usize)
1152 .unwrap_or(0);
1153
1154 let uc = QueryUseCases::new(&self.runtime);
1155 let page = uc
1156 .scan(ScanCollectionInput {
1157 collection: collection.to_string(),
1158 offset,
1159 limit,
1160 })
1161 .map_err(|e| format!("{}", e))?;
1162
1163 let json = crate::presentation::entity_json::scan_page_json(&page);
1164 json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
1165 }
1166
1167 fn tool_graph_centrality(&self, args: &JsonValue) -> Result<String, String> {
1168 let algorithm_str = get_str_field(args, "algorithm")?;
1169 let algo = match algorithm_str {
1170 "degree" => RuntimeGraphCentralityAlgorithm::Degree,
1171 "closeness" => RuntimeGraphCentralityAlgorithm::Closeness,
1172 "betweenness" => RuntimeGraphCentralityAlgorithm::Betweenness,
1173 "eigenvector" => RuntimeGraphCentralityAlgorithm::Eigenvector,
1174 "pagerank" => RuntimeGraphCentralityAlgorithm::PageRank,
1175 _ => return Err(format!("unknown algorithm: {algorithm_str}")),
1176 };
1177
1178 let uc = GraphUseCases::new(&self.runtime);
1179 let result = uc
1180 .centrality(GraphCentralityInput {
1181 algorithm: algo,
1182 top_k: 100,
1183 normalize: true,
1184 max_iterations: None,
1185 epsilon: None,
1186 alpha: None,
1187 projection: None,
1188 })
1189 .map_err(|e| format!("{}", e))?;
1190
1191 let json = crate::presentation::graph_json::graph_centrality_json(&result);
1192 json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
1193 }
1194
1195 fn tool_graph_community(&self, args: &JsonValue) -> Result<String, String> {
1196 let algorithm_str = get_str_field(args, "algorithm")?;
1197 let algo = match algorithm_str {
1198 "label_propagation" => RuntimeGraphCommunityAlgorithm::LabelPropagation,
1199 "louvain" => RuntimeGraphCommunityAlgorithm::Louvain,
1200 _ => return Err(format!("unknown algorithm: {algorithm_str}")),
1201 };
1202 let max_iterations = args
1203 .get("max_iterations")
1204 .and_then(|v| v.as_u64())
1205 .map(|v| v as usize);
1206
1207 let uc = GraphUseCases::new(&self.runtime);
1208 let result = uc
1209 .communities(GraphCommunitiesInput {
1210 algorithm: algo,
1211 min_size: 1,
1212 max_iterations,
1213 resolution: None,
1214 projection: None,
1215 })
1216 .map_err(|e| format!("{}", e))?;
1217
1218 let json = crate::presentation::graph_json::graph_community_json(&result);
1219 json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
1220 }
1221
1222 fn tool_graph_components(&self, args: &JsonValue) -> Result<String, String> {
1223 let mode = match args.get("mode").and_then(|v| v.as_str()) {
1224 Some("strongly_connected") => RuntimeGraphComponentsMode::Strong,
1225 _ => RuntimeGraphComponentsMode::Weak,
1226 };
1227
1228 let uc = GraphUseCases::new(&self.runtime);
1229 let result = uc
1230 .components(GraphComponentsInput {
1231 mode,
1232 min_size: 1,
1233 projection: None,
1234 })
1235 .map_err(|e| format!("{}", e))?;
1236
1237 let json = crate::presentation::graph_json::graph_components_json(&result);
1238 json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
1239 }
1240
1241 fn tool_graph_cycles(&self, args: &JsonValue) -> Result<String, String> {
1242 let max_length = args
1243 .get("max_length")
1244 .and_then(|v| v.as_u64())
1245 .map(|v| v as usize)
1246 .unwrap_or(10);
1247 let max_cycles = args
1248 .get("max_cycles")
1249 .and_then(|v| v.as_u64())
1250 .map(|v| v as usize)
1251 .unwrap_or(100);
1252
1253 let uc = GraphUseCases::new(&self.runtime);
1254 let result = uc
1255 .cycles(GraphCyclesInput {
1256 max_length,
1257 max_cycles,
1258 projection: None,
1259 })
1260 .map_err(|e| format!("{}", e))?;
1261
1262 let json = crate::presentation::graph_json::graph_cycles_json(&result);
1263 json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
1264 }
1265
1266 fn tool_graph_clustering(&self, _args: &JsonValue) -> Result<String, String> {
1267 let uc = GraphUseCases::new(&self.runtime);
1268 let result = uc
1269 .clustering(GraphClusteringInput {
1270 top_k: 100,
1271 include_triangles: true,
1272 projection: None,
1273 })
1274 .map_err(|e| format!("{}", e))?;
1275
1276 let json = crate::presentation::graph_json::graph_clustering_json(&result);
1277 json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
1278 }
1279
1280 fn tool_create_collection(&self, args: &JsonValue) -> Result<String, String> {
1281 let name = get_str_field(args, "name")?;
1282 self.runtime
1283 .db()
1284 .store()
1285 .create_collection(name)
1286 .map_err(|e| format!("{e:?}"))?;
1287 let mut resp = Map::new();
1288 resp.insert("ok".into(), JsonValue::Bool(true));
1289 resp.insert("collection".into(), JsonValue::String(name.to_string()));
1290 json_to_string(&JsonValue::Object(resp)).map_err(|e| format!("serialization error: {}", e))
1291 }
1292
1293 fn tool_drop_collection(&self, args: &JsonValue) -> Result<String, String> {
1294 let name = get_str_field(args, "name")?;
1295 self.runtime
1296 .db()
1297 .store()
1298 .drop_collection(name)
1299 .map_err(|e| format!("{e:?}"))?;
1300 let mut resp = Map::new();
1301 resp.insert("ok".into(), JsonValue::Bool(true));
1302 resp.insert("dropped".into(), JsonValue::String(name.to_string()));
1303 json_to_string(&JsonValue::Object(resp)).map_err(|e| format!("serialization error: {}", e))
1304 }
1305
1306 fn tool_rql_validate(&self, args: &JsonValue) -> Result<String, String> {
1307 let rql = get_str_field(args, "rql")?;
1308 let verdict = reddb_rql::knowledge::validate_rql(rql);
1309 let json = rql_validation_json(&verdict);
1310 json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
1311 }
1312
1313 fn tool_rql_explain(&self, args: &JsonValue) -> Result<String, String> {
1314 let rql = get_str_field(args, "rql")?;
1315 let explanation = reddb_rql::knowledge::explain_rql(rql);
1316
1317 let mut obj = Map::new();
1318 obj.insert(
1319 "validation".to_string(),
1320 rql_validation_json(&explanation.validation),
1321 );
1322 if let Some(ast) = &explanation.ast {
1323 obj.insert("ast".to_string(), JsonValue::String(ast.clone()));
1324 }
1325 if let Some(optimized) = &explanation.optimized_ast {
1326 obj.insert(
1327 "optimized_ast".to_string(),
1328 JsonValue::String(optimized.clone()),
1329 );
1330 }
1331 obj.insert(
1332 "applied_passes".to_string(),
1333 JsonValue::Array(
1334 explanation
1335 .applied_passes
1336 .iter()
1337 .cloned()
1338 .map(JsonValue::String)
1339 .collect(),
1340 ),
1341 );
1342 json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
1343 }
1344
1345 fn tool_type_of(&self, args: &JsonValue) -> Result<String, String> {
1351 let type_name = args.get("type").and_then(|v| v.as_str());
1352 let value = args.get("value");
1353 if type_name.is_none() && value.is_none() {
1354 return Err(
1355 "provide one of 'type' (a type name) or 'value' (a JSON literal)".to_string(),
1356 );
1357 }
1358 let facts =
1359 reddb_types::knowledge::type_of_json(type_name, value).ok_or_else(
1360 || match type_name {
1361 Some(name) => format!("unknown type name '{name}'"),
1362 None => "could not resolve a type from the request".to_string(),
1363 },
1364 )?;
1365 json_to_string(&facts).map_err(|e| format!("serialization error: {}", e))
1366 }
1367
1368 fn tool_explain_connection(&self, args: &JsonValue) -> Result<String, String> {
1369 let uri = get_str_field(args, "uri")?;
1370 let target = reddb_wire::parse(uri).map_err(|err| err.to_string())?;
1371 let json = connection_explanation_json(uri, &target);
1372 json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
1373 }
1374}
1375
1376fn format_mcp_ask_parse_error(err: crate::runtime::ai::mcp_ask_tool::ParseError) -> String {
1381 use crate::runtime::ai::mcp_ask_tool::ParseError;
1382
1383 match err {
1384 ParseError::NotAnObject => "arguments must be an object".to_string(),
1385 ParseError::MissingQuestion => "missing required field 'question'".to_string(),
1386 ParseError::QuestionWrongType => "field 'question' must be a string".to_string(),
1387 ParseError::WrongType { path, expected } => {
1388 format!("{path} must be {expected}")
1389 }
1390 ParseError::OutOfRange { path, detail } => {
1391 format!("{path} out of range: {detail}")
1392 }
1393 ParseError::CacheAndNocache => {
1394 "options.cache and options.nocache are mutually exclusive".to_string()
1395 }
1396 ParseError::UnknownOption { path } => format!("unknown option {path}"),
1397 }
1398}
1399
1400fn connection_explanation_json(uri: &str, target: &reddb_wire::ConnectionTarget) -> JsonValue {
1401 let scheme = connection_uri_scheme(uri);
1402 let mut obj = Map::new();
1403 obj.insert("uri".to_string(), JsonValue::String(uri.to_string()));
1404 obj.insert("scheme".to_string(), JsonValue::String(scheme.clone()));
1405
1406 let mut target_obj = Map::new();
1407 match target {
1408 reddb_wire::ConnectionTarget::Memory => {
1409 obj.insert(
1410 "mode".to_string(),
1411 JsonValue::String("embedded".to_string()),
1412 );
1413 obj.insert(
1414 "transport".to_string(),
1415 JsonValue::String("memory".to_string()),
1416 );
1417 obj.insert(
1418 "topology".to_string(),
1419 JsonValue::String("ephemeral".to_string()),
1420 );
1421 target_obj.insert("kind".to_string(), JsonValue::String("memory".to_string()));
1422 }
1423 reddb_wire::ConnectionTarget::File { path } => {
1424 obj.insert(
1425 "mode".to_string(),
1426 JsonValue::String("embedded".to_string()),
1427 );
1428 obj.insert(
1429 "transport".to_string(),
1430 JsonValue::String("filesystem".to_string()),
1431 );
1432 obj.insert(
1433 "topology".to_string(),
1434 JsonValue::String("persisted".to_string()),
1435 );
1436 target_obj.insert("kind".to_string(), JsonValue::String("file".to_string()));
1437 target_obj.insert(
1438 "path".to_string(),
1439 JsonValue::String(path.display().to_string()),
1440 );
1441 }
1442 reddb_wire::ConnectionTarget::Grpc { endpoint } => {
1443 obj.insert("mode".to_string(), JsonValue::String("remote".to_string()));
1444 obj.insert(
1445 "transport".to_string(),
1446 JsonValue::String(
1447 if scheme == "grpcs" {
1448 "grpc_tls"
1449 } else {
1450 "grpc"
1451 }
1452 .to_string(),
1453 ),
1454 );
1455 obj.insert(
1456 "topology".to_string(),
1457 JsonValue::String("single".to_string()),
1458 );
1459 target_obj.insert("kind".to_string(), JsonValue::String("grpc".to_string()));
1460 target_obj.insert("endpoint".to_string(), JsonValue::String(endpoint.clone()));
1461 }
1462 reddb_wire::ConnectionTarget::GrpcCluster {
1463 primary,
1464 replicas,
1465 force_primary,
1466 } => {
1467 obj.insert("mode".to_string(), JsonValue::String("remote".to_string()));
1468 obj.insert(
1469 "transport".to_string(),
1470 JsonValue::String(cluster_transport_label(&scheme).to_string()),
1471 );
1472 obj.insert(
1473 "topology".to_string(),
1474 JsonValue::String("primary_replica".to_string()),
1475 );
1476 target_obj.insert(
1477 "kind".to_string(),
1478 JsonValue::String("primary_replica".to_string()),
1479 );
1480 target_obj.insert("primary".to_string(), JsonValue::String(primary.clone()));
1481 target_obj.insert(
1482 "replicas".to_string(),
1483 JsonValue::Array(replicas.iter().cloned().map(JsonValue::String).collect()),
1484 );
1485 target_obj.insert("force_primary".to_string(), JsonValue::Bool(*force_primary));
1486 }
1487 reddb_wire::ConnectionTarget::Http { base_url } => {
1488 obj.insert("mode".to_string(), JsonValue::String("remote".to_string()));
1489 obj.insert(
1490 "transport".to_string(),
1491 JsonValue::String(if scheme == "https" { "https" } else { "http" }.to_string()),
1492 );
1493 obj.insert(
1494 "topology".to_string(),
1495 JsonValue::String("single".to_string()),
1496 );
1497 target_obj.insert("kind".to_string(), JsonValue::String("http".to_string()));
1498 target_obj.insert("base_url".to_string(), JsonValue::String(base_url.clone()));
1499 }
1500 reddb_wire::ConnectionTarget::RedWire { host, port, tls } => {
1501 obj.insert("mode".to_string(), JsonValue::String("remote".to_string()));
1502 obj.insert(
1503 "transport".to_string(),
1504 JsonValue::String(if *tls { "redwire_tls" } else { "redwire" }.to_string()),
1505 );
1506 obj.insert(
1507 "topology".to_string(),
1508 JsonValue::String("single".to_string()),
1509 );
1510 target_obj.insert("kind".to_string(), JsonValue::String("redwire".to_string()));
1511 target_obj.insert("host".to_string(), JsonValue::String(host.clone()));
1512 target_obj.insert("port".to_string(), JsonValue::Number(*port as f64));
1513 target_obj.insert("tls".to_string(), JsonValue::Bool(*tls));
1514 }
1515 reddb_wire::ConnectionTarget::WsNative { host, port, tls } => {
1516 obj.insert("mode".to_string(), JsonValue::String("remote".to_string()));
1517 obj.insert(
1518 "transport".to_string(),
1519 JsonValue::String(if *tls { "redwire_wss" } else { "redwire_ws" }.to_string()),
1520 );
1521 obj.insert(
1522 "topology".to_string(),
1523 JsonValue::String("single".to_string()),
1524 );
1525 target_obj.insert(
1526 "kind".to_string(),
1527 JsonValue::String("redwire_websocket".to_string()),
1528 );
1529 target_obj.insert("host".to_string(), JsonValue::String(host.clone()));
1530 target_obj.insert("port".to_string(), JsonValue::Number(*port as f64));
1531 target_obj.insert("tls".to_string(), JsonValue::Bool(*tls));
1532 }
1533 }
1534 obj.insert("target".to_string(), JsonValue::Object(target_obj));
1535 JsonValue::Object(obj)
1536}
1537
1538fn connection_uri_scheme(uri: &str) -> String {
1539 uri.split_once(':')
1540 .map(|(scheme, _)| scheme.to_ascii_lowercase())
1541 .unwrap_or_default()
1542}
1543
1544fn cluster_transport_label(scheme: &str) -> &'static str {
1545 match scheme {
1546 "red" => "redwire",
1547 "reds" => "redwire_tls",
1548 "grpcs" => "grpc_tls",
1549 _ => "grpc",
1550 }
1551}
1552
1553fn rql_validation_json(verdict: &reddb_rql::knowledge::RqlValidation) -> JsonValue {
1556 let mut obj = Map::new();
1557 obj.insert("valid".to_string(), JsonValue::Bool(verdict.valid));
1558 if let Some(statement) = &verdict.statement {
1559 obj.insert(
1560 "statement".to_string(),
1561 JsonValue::String(statement.clone()),
1562 );
1563 }
1564 obj.insert(
1565 "has_with_clause".to_string(),
1566 JsonValue::Bool(verdict.has_with_clause),
1567 );
1568 if let Some(error) = &verdict.error {
1569 obj.insert("error".to_string(), rql_diagnostic_json(error));
1570 }
1571 JsonValue::Object(obj)
1572}
1573
1574fn rql_diagnostic_json(diag: &reddb_rql::knowledge::RqlDiagnostic) -> JsonValue {
1576 let mut obj = Map::new();
1577 obj.insert(
1578 "message".to_string(),
1579 JsonValue::String(diag.message.clone()),
1580 );
1581 obj.insert("line".to_string(), JsonValue::Number(diag.line as f64));
1582 obj.insert("column".to_string(), JsonValue::Number(diag.column as f64));
1583 obj.insert("offset".to_string(), JsonValue::Number(diag.offset as f64));
1584 obj.insert("kind".to_string(), JsonValue::String(diag.kind.clone()));
1585 obj.insert(
1586 "expected".to_string(),
1587 JsonValue::Array(
1588 diag.expected
1589 .iter()
1590 .cloned()
1591 .map(JsonValue::String)
1592 .collect(),
1593 ),
1594 );
1595 JsonValue::Object(obj)
1596}
1597
1598fn parse_direction(s: Option<&str>) -> RuntimeGraphDirection {
1599 match s {
1600 Some("incoming") => RuntimeGraphDirection::Incoming,
1601 Some("both") => RuntimeGraphDirection::Both,
1602 _ => RuntimeGraphDirection::Outgoing,
1603 }
1604}
1605
1606fn parse_metadata_arg(
1608 args: &JsonValue,
1609) -> Result<Vec<(String, crate::storage::unified::MetadataValue)>, String> {
1610 match args.get("metadata").and_then(|v| v.as_object()) {
1611 Some(obj) => {
1612 let mut out = Vec::with_capacity(obj.len());
1613 for (key, value) in obj {
1614 let mv = crate::application::entity::json_to_metadata_value(value)
1615 .map_err(|e| format!("{}", e))?;
1616 out.push((key.clone(), mv));
1617 }
1618 Ok(out)
1619 }
1620 None => Ok(vec![]),
1621 }
1622}
1623
1624fn parse_string_array_arg(args: &JsonValue, field: &str) -> Result<Vec<String>, String> {
1625 match args.get(field) {
1626 None | Some(JsonValue::Null) => Ok(Vec::new()),
1627 Some(JsonValue::Array(values)) => values
1628 .iter()
1629 .map(|value| {
1630 value
1631 .as_str()
1632 .map(ToOwned::to_owned)
1633 .ok_or_else(|| format!("field '{field}' must be an array of strings"))
1634 })
1635 .collect(),
1636 _ => Err(format!("field '{field}' must be an array of strings")),
1637 }
1638}
1639
1640fn mcp_keyed_ident(value: &str) -> Result<String, String> {
1641 if !value.is_empty()
1642 && value
1643 .bytes()
1644 .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'.')
1645 {
1646 Ok(value.to_string())
1647 } else {
1648 Err(
1649 "keyed collection and key names must use letters, numbers, underscores, or dots"
1650 .to_string(),
1651 )
1652 }
1653}
1654
1655fn mcp_value_literal(value: &JsonValue) -> Result<String, String> {
1656 match value {
1657 JsonValue::String(value) => Ok(format!("'{}'", value.replace('\'', "''"))),
1658 JsonValue::Number(value) => Ok(value.to_string()),
1659 JsonValue::Bool(value) => Ok(value.to_string()),
1660 JsonValue::Null => Ok("NULL".to_string()),
1661 JsonValue::Array(_) | JsonValue::Object(_) => {
1662 json_to_string(value).map_err(|err| format!("serialization error: {err}"))
1663 }
1664 }
1665}
1666
1667fn append_mcp_tags_clause(sql: &mut String, tags: &[String]) {
1668 if tags.is_empty() {
1669 return;
1670 }
1671 sql.push_str(" TAGS [");
1672 for (index, tag) in tags.iter().enumerate() {
1673 if index > 0 {
1674 sql.push_str(", ");
1675 }
1676 sql.push('\'');
1677 sql.push_str(&tag.replace('\'', "''"));
1678 sql.push('\'');
1679 }
1680 sql.push(']');
1681}
1682
1683fn reject_mcp_volatile_options(args: &JsonValue, domain: &str) -> Result<(), String> {
1684 for field in ["ttl", "ttl_ms", "expire", "expire_ms", "expires_at"] {
1685 if args.get(field).is_some() {
1686 return Err(format!(
1687 "{domain} does not support TTL or expiration options"
1688 ));
1689 }
1690 }
1691 Ok(())
1692}
1693
1694fn get_str_field<'a>(args: &'a JsonValue, field: &str) -> Result<&'a str, String> {
1696 args.get(field)
1697 .and_then(|v| v.as_str())
1698 .ok_or_else(|| format!("missing '{field}'"))
1699}
1700
1701impl McpServer {
1703 fn tool_auth_bootstrap(&self, args: &JsonValue) -> Result<String, String> {
1704 let username = get_str_field(args, "username")?;
1705 let password = get_str_field(args, "password")?;
1706
1707 let br = self
1708 .auth_store
1709 .bootstrap(username, password)
1710 .map_err(|e| e.to_string())?;
1711
1712 let mut result = Map::new();
1713 result.insert("ok".into(), JsonValue::Bool(true));
1714 result.insert("username".into(), JsonValue::String(br.user.username));
1715 result.insert(
1716 "role".into(),
1717 JsonValue::String(br.user.role.as_str().into()),
1718 );
1719 result.insert("api_key".into(), JsonValue::String(br.api_key.key));
1720 result.insert("api_key_name".into(), JsonValue::String(br.api_key.name));
1721 if let Some(cert) = br.certificate {
1722 result.insert("certificate".into(), JsonValue::String(cert));
1723 result.insert(
1724 "message".into(),
1725 JsonValue::String(
1726 "Save this certificate — it is the ONLY way to unseal the vault after restart."
1727 .into(),
1728 ),
1729 );
1730 } else {
1731 result.insert(
1732 "message".into(),
1733 JsonValue::String(
1734 "First admin user created. Save the API key — it won't be shown again.".into(),
1735 ),
1736 );
1737 }
1738 json_to_string(&JsonValue::Object(result))
1739 }
1740
1741 fn tool_auth_create_user(&self, args: &JsonValue) -> Result<String, String> {
1742 let username = get_str_field(args, "username")?;
1743 let password = get_str_field(args, "password")?;
1744 let role_str = get_str_field(args, "role")?;
1745 let role = Role::from_str(role_str).ok_or_else(|| format!("invalid role: {role_str}"))?;
1746
1747 self.auth_store
1748 .create_user(username, password, role)
1749 .map_err(|e| e.to_string())?;
1750
1751 let mut result = Map::new();
1752 result.insert("ok".into(), JsonValue::Bool(true));
1753 result.insert("username".into(), JsonValue::String(username.into()));
1754 result.insert("role".into(), JsonValue::String(role.as_str().into()));
1755 json_to_string(&JsonValue::Object(result))
1756 }
1757
1758 fn tool_auth_login(&self, args: &JsonValue) -> Result<String, String> {
1759 let username = get_str_field(args, "username")?;
1760 let password = get_str_field(args, "password")?;
1761
1762 let session = self
1763 .auth_store
1764 .authenticate(username, password)
1765 .map_err(|e| e.to_string())?;
1766
1767 let mut result = Map::new();
1768 result.insert("ok".into(), JsonValue::Bool(true));
1769 result.insert("token".into(), JsonValue::String(session.token));
1770 result.insert("username".into(), JsonValue::String(session.username));
1771 result.insert(
1772 "role".into(),
1773 JsonValue::String(session.role.as_str().into()),
1774 );
1775 result.insert(
1776 "expires_at".into(),
1777 JsonValue::Number(session.expires_at as f64),
1778 );
1779 json_to_string(&JsonValue::Object(result))
1780 }
1781
1782 fn tool_auth_create_api_key(&self, args: &JsonValue) -> Result<String, String> {
1783 let username = get_str_field(args, "username")?;
1784 let name = get_str_field(args, "name")?;
1785 let role_str = get_str_field(args, "role")?;
1786 let role = Role::from_str(role_str).ok_or_else(|| format!("invalid role: {role_str}"))?;
1787
1788 let key = self
1789 .auth_store
1790 .create_api_key(username, name, role)
1791 .map_err(|e| e.to_string())?;
1792
1793 let mut result = Map::new();
1794 result.insert("ok".into(), JsonValue::Bool(true));
1795 result.insert("key".into(), JsonValue::String(key.key));
1796 result.insert("name".into(), JsonValue::String(key.name));
1797 result.insert("role".into(), JsonValue::String(key.role.as_str().into()));
1798 json_to_string(&JsonValue::Object(result))
1799 }
1800
1801 fn tool_auth_list_users(&self) -> Result<String, String> {
1802 let users = self.auth_store.list_users();
1803 let arr: Vec<JsonValue> = users
1804 .into_iter()
1805 .map(|u| {
1806 let mut obj = Map::new();
1807 obj.insert("username".into(), JsonValue::String(u.username));
1808 obj.insert("role".into(), JsonValue::String(u.role.as_str().into()));
1809 obj.insert("enabled".into(), JsonValue::Bool(u.enabled));
1810 obj.insert(
1811 "api_key_count".into(),
1812 JsonValue::Number(u.api_keys.len() as f64),
1813 );
1814 JsonValue::Object(obj)
1815 })
1816 .collect();
1817 json_to_string(&JsonValue::Array(arr))
1818 }
1819}
1820
1821#[cfg(test)]
1822mod tests {
1823 use super::*;
1824 use std::io::{Read, Write};
1825 use std::net::{SocketAddr, TcpListener, TcpStream};
1826 use std::sync::atomic::{AtomicBool, Ordering};
1827 use std::sync::{Arc, Mutex};
1828 use std::thread::{self, JoinHandle};
1829 use std::time::Duration;
1830
1831 static ASK_ENV_LOCK: Mutex<()> = Mutex::new(());
1832
1833 fn make_server() -> McpServer {
1834 let rt = RedDBRuntime::in_memory().expect("in-memory runtime");
1835 McpServer::new(rt)
1836 }
1837
1838 fn parse_json(s: &str) -> JsonValue {
1839 json_from_str(s).expect("valid json")
1840 }
1841
1842 #[test]
1843 fn initialize_advertises_resources_capability() {
1844 let mut srv = make_server();
1845 let response = srv.handle_initialize(Some(&JsonValue::Number(1.0)));
1846 let parsed = parse_json(&response);
1847 let resources = parsed
1848 .get("result")
1849 .and_then(|result| result.get("capabilities"))
1850 .and_then(|caps| caps.get("resources"))
1851 .expect("resources capability advertised");
1852 assert_eq!(
1853 resources.get("subscribe").and_then(JsonValue::as_bool),
1854 Some(false)
1855 );
1856 }
1857
1858 #[test]
1859 fn resources_list_includes_rql_knowledge() {
1860 let mut srv = make_server();
1861 let request = parse_json(r#"{"jsonrpc":"2.0","id":7,"method":"resources/list"}"#);
1862 let response = srv.handle_message(&request).expect("response");
1863 let parsed = parse_json(&response);
1864 let resources = parsed
1865 .get("result")
1866 .and_then(|result| result.get("resources"))
1867 .and_then(JsonValue::as_array)
1868 .expect("resources array");
1869
1870 let rql = resources
1871 .iter()
1872 .find(|res| res.get("uri").and_then(JsonValue::as_str) == Some("reddb://knowledge/rql"))
1873 .expect("rql knowledge resource listed");
1874 assert_eq!(
1875 rql.get("mimeType").and_then(JsonValue::as_str),
1876 Some("text/markdown")
1877 );
1878 assert!(rql.get("name").and_then(JsonValue::as_str).is_some());
1879 }
1880
1881 #[test]
1882 fn resources_list_includes_connection_knowledge() {
1883 let mut srv = make_server();
1884 let request = parse_json(r#"{"jsonrpc":"2.0","id":11,"method":"resources/list"}"#);
1885 let response = srv.handle_message(&request).expect("response");
1886 let parsed = parse_json(&response);
1887 let resources = parsed
1888 .get("result")
1889 .and_then(|result| result.get("resources"))
1890 .and_then(JsonValue::as_array)
1891 .expect("resources array");
1892
1893 let connection = resources
1894 .iter()
1895 .find(|res| {
1896 res.get("uri").and_then(JsonValue::as_str)
1897 == Some(reddb_wire::knowledge::RESOURCE_URI)
1898 })
1899 .expect("connection knowledge resource listed");
1900 assert_eq!(
1901 connection.get("mimeType").and_then(JsonValue::as_str),
1902 Some("text/markdown")
1903 );
1904 assert!(connection.get("name").and_then(JsonValue::as_str).is_some());
1905 }
1906
1907 #[test]
1908 fn resources_read_returns_generated_rql_reference() {
1909 let mut srv = make_server();
1910 let request = parse_json(
1911 r#"{"jsonrpc":"2.0","id":8,"method":"resources/read","params":{"uri":"reddb://knowledge/rql"}}"#,
1912 );
1913 let response = srv.handle_message(&request).expect("response");
1914 let parsed = parse_json(&response);
1915 let text = parsed
1916 .get("result")
1917 .and_then(|result| result.get("contents"))
1918 .and_then(JsonValue::as_array)
1919 .and_then(|contents| contents.first())
1920 .and_then(|item| item.get("text"))
1921 .and_then(JsonValue::as_str)
1922 .expect("resource text");
1923
1924 assert_eq!(text, reddb_rql::knowledge::rql_reference_markdown());
1927 assert!(text.contains("`SELECT`"), "keyword missing: {text:.120}");
1930 assert!(text.contains("`COUNT`"), "function missing");
1931 }
1932
1933 #[test]
1934 fn resources_read_returns_generated_type_reference() {
1935 let mut srv = make_server();
1936 let request = parse_json(
1937 r#"{"jsonrpc":"2.0","id":10,"method":"resources/read","params":{"uri":"reddb://knowledge/types"}}"#,
1938 );
1939 let response = srv.handle_message(&request).expect("response");
1940 let parsed = parse_json(&response);
1941 let text = parsed
1942 .get("result")
1943 .and_then(|result| result.get("contents"))
1944 .and_then(JsonValue::as_array)
1945 .and_then(|contents| contents.first())
1946 .and_then(|item| item.get("text"))
1947 .and_then(JsonValue::as_str)
1948 .expect("resource text");
1949
1950 assert_eq!(text, reddb_types::knowledge::type_reference_markdown());
1953 assert!(
1956 text.contains("`INTEGER`"),
1957 "value type missing: {text:.120}"
1958 );
1959 assert!(text.contains("Multi-model map"), "multi-model map missing");
1960 }
1961
1962 #[test]
1963 fn resources_read_returns_generated_connection_reference() {
1964 let mut srv = make_server();
1965 let request = parse_json(
1966 r#"{"jsonrpc":"2.0","id":12,"method":"resources/read","params":{"uri":"reddb://knowledge/connections"}}"#,
1967 );
1968 let response = srv.handle_message(&request).expect("response");
1969 let parsed = parse_json(&response);
1970 let text = parsed
1971 .get("result")
1972 .and_then(|result| result.get("contents"))
1973 .and_then(JsonValue::as_array)
1974 .and_then(|contents| contents.first())
1975 .and_then(|item| item.get("text"))
1976 .and_then(JsonValue::as_str)
1977 .expect("resource text");
1978
1979 assert_eq!(text, reddb_wire::knowledge::connection_reference_markdown());
1982 assert!(text.contains("`red://`"), "red scheme missing: {text:.120}");
1983 assert!(text.contains("`reds://`"), "reds scheme missing");
1984 assert!(
1985 text.contains("principal transport"),
1986 "principal transport narrative missing"
1987 );
1988 assert!(
1989 text.contains("Embedded / standalone"),
1990 "embedded topology missing"
1991 );
1992 assert!(text.contains("Serverless"), "serverless topology missing");
1993 assert!(
1994 text.contains("Primary-replica"),
1995 "primary-replica topology missing"
1996 );
1997 assert!(text.contains("Cluster"), "cluster topology missing");
1998 }
1999
2000 #[test]
2001 fn resources_read_unknown_uri_errors() {
2002 let mut srv = make_server();
2003 let request = parse_json(
2004 r#"{"jsonrpc":"2.0","id":9,"method":"resources/read","params":{"uri":"reddb://knowledge/nope"}}"#,
2005 );
2006 let response = srv.handle_message(&request).expect("response");
2007 let parsed = parse_json(&response);
2008 assert!(
2009 parsed.get("error").is_some(),
2010 "unknown resource must error: {response}"
2011 );
2012 }
2013
2014 #[test]
2015 fn tools_list_registers_reddb_ask_descriptor() {
2016 let srv = make_server();
2017 let response = srv.handle_tools_list(Some(&JsonValue::Number(1.0)));
2018 let parsed = parse_json(&response);
2019 let tools = parsed
2020 .get("result")
2021 .and_then(|result| result.get("tools"))
2022 .and_then(JsonValue::as_array)
2023 .expect("tools array");
2024
2025 let ask = tools
2026 .iter()
2027 .find(|tool| tool.get("name").and_then(JsonValue::as_str) == Some("reddb.ask"))
2028 .expect("reddb.ask registered");
2029
2030 let desc = ask
2031 .get("description")
2032 .and_then(JsonValue::as_str)
2033 .expect("description");
2034 assert!(desc.contains("citations"), "description: {desc}");
2035 assert!(desc.contains("sources_flat"), "description: {desc}");
2036 assert!(desc.contains("URN"), "description: {desc}");
2037
2038 let options = ask
2039 .get("inputSchema")
2040 .and_then(|schema| schema.get("properties"))
2041 .and_then(|props| props.get("options"))
2042 .and_then(|opts| opts.get("properties"))
2043 .and_then(JsonValue::as_object)
2044 .expect("options properties");
2045 for key in [
2046 "strict",
2047 "using",
2048 "model",
2049 "limit",
2050 "min_score",
2051 "depth",
2052 "temperature",
2053 "seed",
2054 "cache",
2055 "nocache",
2056 ] {
2057 assert!(
2058 options.contains_key(key),
2059 "missing option {key} in {options:?}"
2060 );
2061 }
2062 }
2063
2064 #[test]
2065 fn tools_call_reddb_ask_uses_typed_argument_parser() {
2066 let srv = make_server();
2067 let params = parse_json(
2068 r#"{
2069 "name": "reddb.ask",
2070 "arguments": {
2071 "question": "what cites this?",
2072 "options": { "tempurature": 0.2 }
2073 }
2074 }"#,
2075 );
2076
2077 let response = srv.handle_tools_call(Some(&JsonValue::Number(1.0)), Some(¶ms));
2078 let parsed = parse_json(&response);
2079 let result = parsed.get("result").expect("result");
2080 assert_eq!(
2081 result.get("isError").and_then(JsonValue::as_bool),
2082 Some(true)
2083 );
2084 let text = result
2085 .get("content")
2086 .and_then(JsonValue::as_array)
2087 .and_then(|content| content.first())
2088 .and_then(|item| item.get("text"))
2089 .and_then(JsonValue::as_str)
2090 .expect("error text");
2091 assert!(text.contains("options.tempurature"), "text: {text}");
2092 }
2093
2094 fn call_type_of(srv: &mut McpServer, arguments: &str) -> JsonValue {
2098 let request = parse_json(&format!(
2099 r#"{{"jsonrpc":"2.0","id":21,"method":"tools/call","params":{{"name":"reddb_type_of","arguments":{arguments}}}}}"#
2100 ));
2101 let response = srv.handle_message(&request).expect("response");
2102 parse_json(&response)
2103 .get("result")
2104 .cloned()
2105 .expect("result object")
2106 }
2107
2108 fn type_of_result_json(result: &JsonValue) -> JsonValue {
2110 assert_ne!(
2111 result.get("isError").and_then(JsonValue::as_bool),
2112 Some(true),
2113 "tool returned an error: {result:?}"
2114 );
2115 let text = result
2116 .get("content")
2117 .and_then(JsonValue::as_array)
2118 .and_then(|content| content.first())
2119 .and_then(|item| item.get("text"))
2120 .and_then(JsonValue::as_str)
2121 .expect("result text");
2122 parse_json(text)
2123 }
2124
2125 #[test]
2126 fn tools_list_includes_reddb_type_of() {
2127 let srv = make_server();
2128 let response = srv.handle_tools_list(Some(&JsonValue::Number(1.0)));
2129 let parsed = parse_json(&response);
2130 let tools = parsed
2131 .get("result")
2132 .and_then(|result| result.get("tools"))
2133 .and_then(JsonValue::as_array)
2134 .expect("tools array");
2135 assert!(
2136 tools
2137 .iter()
2138 .any(|tool| tool.get("name").and_then(JsonValue::as_str) == Some("reddb_type_of")),
2139 "reddb_type_of must be advertised in tools/list"
2140 );
2141 }
2142
2143 #[test]
2144 fn tools_call_type_of_resolves_type_name_from_catalog() {
2145 let mut srv = make_server();
2146 let result = call_type_of(&mut srv, r#"{"type":"int"}"#);
2147 let facts = type_of_result_json(&result);
2148
2149 assert_eq!(
2151 facts.get("canonical_type").and_then(JsonValue::as_str),
2152 Some("INTEGER")
2153 );
2154 assert_eq!(
2155 facts.get("category").and_then(JsonValue::as_str),
2156 Some("Numeric")
2157 );
2158 let casts = facts.get("casts").and_then(JsonValue::as_array).unwrap();
2160 assert!(
2161 casts.iter().any(|c| {
2162 c.get("target").and_then(JsonValue::as_str) == Some("FLOAT")
2163 && c.get("context").and_then(JsonValue::as_str) == Some("implicit")
2164 }),
2165 "casts: {casts:?}"
2166 );
2167 let operators = facts
2169 .get("operators")
2170 .and_then(JsonValue::as_array)
2171 .unwrap();
2172 assert!(
2173 operators
2174 .iter()
2175 .any(|o| o.get("symbol").and_then(JsonValue::as_str) == Some("+")),
2176 "operators: {operators:?}"
2177 );
2178 }
2179
2180 #[test]
2181 fn tools_call_type_of_infers_type_from_literal_value() {
2182 let mut srv = make_server();
2183 let result = call_type_of(&mut srv, r#"{"value":true}"#);
2184 let facts = type_of_result_json(&result);
2185 assert_eq!(
2186 facts.get("canonical_type").and_then(JsonValue::as_str),
2187 Some("BOOLEAN")
2188 );
2189 }
2190
2191 #[test]
2192 fn tools_call_type_of_unknown_name_is_error() {
2193 let mut srv = make_server();
2194 let result = call_type_of(&mut srv, r#"{"type":"frobnicate"}"#);
2195 assert_eq!(
2196 result.get("isError").and_then(JsonValue::as_bool),
2197 Some(true)
2198 );
2199 let text = result
2200 .get("content")
2201 .and_then(JsonValue::as_array)
2202 .and_then(|content| content.first())
2203 .and_then(|item| item.get("text"))
2204 .and_then(JsonValue::as_str)
2205 .expect("error text");
2206 assert!(text.contains("frobnicate"), "text: {text}");
2207 }
2208
2209 #[test]
2210 fn tools_call_type_of_requires_an_argument() {
2211 let mut srv = make_server();
2212 let result = call_type_of(&mut srv, r#"{}"#);
2213 assert_eq!(
2214 result.get("isError").and_then(JsonValue::as_bool),
2215 Some(true)
2216 );
2217 }
2218
2219 #[test]
2220 fn tools_call_reddb_ask_returns_canonical_citation_envelope() {
2221 let _guard = ASK_ENV_LOCK.lock().expect("env lock");
2222 let stub = AskStub::start();
2223 let _api_base = EnvVarGuard::set(
2224 "REDDB_OLLAMA_API_BASE",
2225 &format!("http://{}/v1", stub.addr()),
2226 );
2227
2228 let srv = make_server();
2229 srv.tool_query(&parse_json(
2230 r#"{"sql":"CREATE TABLE travel (id INTEGER, passport TEXT, notes TEXT)"}"#,
2231 ))
2232 .expect("ddl ok");
2233 srv.tool_query(&parse_json(
2234 r#"{"sql":"INSERT INTO travel (id, passport, notes) VALUES (1, 'PT-002', 'incident FDD-12313 escalated')"}"#,
2235 ))
2236 .expect("insert ok");
2237
2238 let params = parse_json(
2239 r#"{
2240 "name": "reddb.ask",
2241 "arguments": {
2242 "question": "passport FDD-12313",
2243 "options": {
2244 "strict": false,
2245 "using": "ollama",
2246 "model": "mock-ask",
2247 "limit": 1,
2248 "min_score": 0,
2249 "depth": 0,
2250 "temperature": 0,
2251 "seed": 0,
2252 "cache": { "ttl": "5m" }
2253 }
2254 }
2255 }"#,
2256 );
2257
2258 let response = srv.handle_tools_call(Some(&JsonValue::Number(1.0)), Some(¶ms));
2259 let parsed = parse_json(&response);
2260 let result = parsed.get("result").expect("result");
2261 assert_ne!(
2262 result.get("isError").and_then(JsonValue::as_bool),
2263 Some(true),
2264 "response: {response}"
2265 );
2266 let text = result
2267 .get("content")
2268 .and_then(JsonValue::as_array)
2269 .and_then(|content| content.first())
2270 .and_then(|item| item.get("text"))
2271 .and_then(JsonValue::as_str)
2272 .expect("tool text");
2273 let envelope = parse_json(text);
2274
2275 assert_eq!(
2276 envelope.get("answer").and_then(JsonValue::as_str),
2277 Some("FDD-12313 escalated [^1].")
2278 );
2279 assert_eq!(
2280 envelope.get("provider").and_then(JsonValue::as_str),
2281 Some("ollama")
2282 );
2283 assert_eq!(
2284 envelope.get("model").and_then(JsonValue::as_str),
2285 Some("mock-ask")
2286 );
2287 assert_eq!(
2288 envelope.get("cache_hit").and_then(JsonValue::as_bool),
2289 Some(false)
2290 );
2291 assert!(envelope
2292 .get("sources_flat")
2293 .and_then(JsonValue::as_array)
2294 .is_some());
2295 assert!(envelope
2296 .get("citations")
2297 .and_then(JsonValue::as_array)
2298 .is_some());
2299 assert!(envelope
2300 .get("validation")
2301 .and_then(JsonValue::as_object)
2302 .is_some());
2303 assert!(
2304 envelope.get("rows").is_none(),
2305 "ASK must not be row-wrapped: {text}"
2306 );
2307 }
2308
2309 #[test]
2310 fn tool_query_without_params_keeps_legacy_path() {
2311 let srv = make_server();
2312 let args = parse_json(r#"{"sql":"SELECT 1 AS one"}"#);
2313 let out = srv.tool_query(&args).expect("query ok");
2314 assert!(out.contains("\"one\""), "expected 'one' column in {out}");
2315 }
2316
2317 #[test]
2318 fn tool_query_binds_int_and_text_params() {
2319 let srv = make_server();
2320 srv.tool_query(&parse_json(
2321 r#"{"sql":"CREATE TABLE mcpp (id INTEGER, name TEXT)"}"#,
2322 ))
2323 .expect("ddl ok");
2324 srv.tool_query(&parse_json(
2325 r#"{"sql":"INSERT INTO mcpp (id, name) VALUES (1, 'Alice')"}"#,
2326 ))
2327 .expect("insert 1");
2328 srv.tool_query(&parse_json(
2329 r#"{"sql":"INSERT INTO mcpp (id, name) VALUES (2, 'Bob')"}"#,
2330 ))
2331 .expect("insert 2");
2332
2333 let out = srv
2334 .tool_query(&parse_json(
2335 r#"{"sql":"SELECT * FROM mcpp WHERE id = $1 AND name = $2","params":[1,"Alice"]}"#,
2336 ))
2337 .expect("query with params ok");
2338 assert!(out.contains("Alice"), "expected Alice in {out}");
2339 assert!(!out.contains("Bob"), "Bob must not match: {out}");
2340 }
2341
2342 #[test]
2343 fn tool_query_params_must_be_array() {
2344 let srv = make_server();
2345 let err = srv
2346 .tool_query(&parse_json(
2347 r#"{"sql":"SELECT 1","params":{"not":"array"}}"#,
2348 ))
2349 .expect_err("must reject non-array params");
2350 assert!(err.contains("array"), "got {err}");
2351 }
2352
2353 #[test]
2354 fn tool_query_param_arity_mismatch_surfaces_error() {
2355 let srv = make_server();
2356 srv.tool_query(&parse_json(r#"{"sql":"CREATE TABLE mcpa (id INTEGER)"}"#))
2357 .expect("ddl ok");
2358 let err = srv
2359 .tool_query(&parse_json(
2360 r#"{"sql":"SELECT * FROM mcpa WHERE id = $1","params":[1,2]}"#,
2361 ))
2362 .expect_err("arity mismatch");
2363 assert!(
2364 err.contains("number of parameters") || err.contains("expects"),
2365 "got {err}"
2366 );
2367 }
2368
2369 #[test]
2370 fn tool_query_vector_param_binds_into_search_similar() {
2371 let srv = make_server();
2372 let out = srv.tool_query(&parse_json(
2373 r#"{"sql":"SEARCH SIMILAR $1 COLLECTION mcpv LIMIT 5","params":[[0.1,0.2,0.3]]}"#,
2374 ));
2375 if let Err(e) = out {
2379 assert!(
2380 !e.contains("placeholder") && !e.contains("Parameter"),
2381 "param did not bind: {e}"
2382 );
2383 }
2384 }
2385
2386 fn call_tool(srv: &McpServer, name: &str, rql: &str) -> JsonValue {
2392 let params = parse_json(
2393 &json_to_string(&{
2394 let mut p = Map::new();
2395 p.insert("name".to_string(), JsonValue::String(name.to_string()));
2396 let mut args = Map::new();
2397 args.insert("rql".to_string(), JsonValue::String(rql.to_string()));
2398 p.insert("arguments".to_string(), JsonValue::Object(args));
2399 JsonValue::Object(p)
2400 })
2401 .expect("serialize params"),
2402 );
2403 let response = srv.handle_tools_call(Some(&JsonValue::Number(1.0)), Some(¶ms));
2404 parse_json(&response)
2405 }
2406
2407 fn tool_result_text(parsed: &JsonValue) -> String {
2408 parsed
2409 .get("result")
2410 .and_then(|r| r.get("content"))
2411 .and_then(JsonValue::as_array)
2412 .and_then(|content| content.first())
2413 .and_then(|item| item.get("text"))
2414 .and_then(JsonValue::as_str)
2415 .expect("tool text")
2416 .to_string()
2417 }
2418
2419 #[test]
2420 fn tools_list_registers_rql_validate_and_explain() {
2421 let srv = make_server();
2422 let response = srv.handle_tools_list(Some(&JsonValue::Number(1.0)));
2423 let parsed = parse_json(&response);
2424 let tools = parsed
2425 .get("result")
2426 .and_then(|result| result.get("tools"))
2427 .and_then(JsonValue::as_array)
2428 .expect("tools array");
2429 let names: Vec<&str> = tools
2430 .iter()
2431 .filter_map(|t| t.get("name").and_then(JsonValue::as_str))
2432 .collect();
2433 assert!(names.contains(&"reddb_rql_validate"), "{names:?}");
2434 assert!(names.contains(&"reddb_rql_explain"), "{names:?}");
2435 }
2436
2437 #[test]
2438 fn tools_call_rql_validate_accepts_valid_query() {
2439 let srv = make_server();
2440 let parsed = call_tool(&srv, "reddb_rql_validate", "SELECT * FROM users");
2441 let result = parsed.get("result").expect("result");
2442 assert_ne!(
2443 result.get("isError").and_then(JsonValue::as_bool),
2444 Some(true)
2445 );
2446 let envelope = parse_json(&tool_result_text(&parsed));
2447 assert_eq!(
2448 envelope.get("valid").and_then(JsonValue::as_bool),
2449 Some(true)
2450 );
2451 assert_eq!(
2452 envelope.get("statement").and_then(JsonValue::as_str),
2453 Some("Table")
2454 );
2455 }
2456
2457 #[test]
2458 fn tools_call_rql_validate_reports_structured_error() {
2459 let srv = make_server();
2460 let parsed = call_tool(&srv, "reddb_rql_validate", "SELECT * FROM");
2461 let envelope = parse_json(&tool_result_text(&parsed));
2462 assert_eq!(
2463 envelope.get("valid").and_then(JsonValue::as_bool),
2464 Some(false)
2465 );
2466 let error = envelope.get("error").expect("structured error");
2467 assert!(error
2468 .get("message")
2469 .and_then(JsonValue::as_str)
2470 .is_some_and(|m| !m.is_empty()));
2471 assert!(error.get("line").and_then(JsonValue::as_f64).is_some());
2472 assert!(error
2473 .get("kind")
2474 .and_then(JsonValue::as_str)
2475 .is_some_and(|k| !k.is_empty()));
2476 }
2477
2478 #[test]
2479 fn tools_call_rql_explain_returns_ast_and_plan() {
2480 let srv = make_server();
2481 let parsed = call_tool(
2482 &srv,
2483 "reddb_rql_explain",
2484 "SELECT * FROM users WHERE id = 1",
2485 );
2486 let envelope = parse_json(&tool_result_text(&parsed));
2487 assert_eq!(
2488 envelope
2489 .get("validation")
2490 .and_then(|v| v.get("valid"))
2491 .and_then(JsonValue::as_bool),
2492 Some(true)
2493 );
2494 assert!(envelope.get("ast").and_then(JsonValue::as_str).is_some());
2495 assert!(envelope
2496 .get("optimized_ast")
2497 .and_then(JsonValue::as_str)
2498 .is_some());
2499 assert!(envelope
2500 .get("applied_passes")
2501 .and_then(JsonValue::as_array)
2502 .is_some());
2503 }
2504
2505 #[test]
2506 fn tools_call_rql_validate_missing_field_errors() {
2507 let srv = make_server();
2508 let params = parse_json(r#"{"name":"reddb_rql_validate","arguments":{}}"#);
2509 let response = srv.handle_tools_call(Some(&JsonValue::Number(1.0)), Some(¶ms));
2510 let parsed = parse_json(&response);
2511 let result = parsed.get("result").expect("result");
2512 assert_eq!(
2513 result.get("isError").and_then(JsonValue::as_bool),
2514 Some(true)
2515 );
2516 }
2517
2518 fn call_explain_connection(srv: &mut McpServer, uri: &str) -> JsonValue {
2519 let request = parse_json(
2520 &json_to_string(&{
2521 let mut request = Map::new();
2522 request.insert("jsonrpc".to_string(), JsonValue::String("2.0".to_string()));
2523 request.insert("id".to_string(), JsonValue::Number(31.0));
2524 request.insert(
2525 "method".to_string(),
2526 JsonValue::String("tools/call".to_string()),
2527 );
2528
2529 let mut args = Map::new();
2530 args.insert("uri".to_string(), JsonValue::String(uri.to_string()));
2531
2532 let mut params = Map::new();
2533 params.insert(
2534 "name".to_string(),
2535 JsonValue::String("reddb_explain_connection".to_string()),
2536 );
2537 params.insert("arguments".to_string(), JsonValue::Object(args));
2538 request.insert("params".to_string(), JsonValue::Object(params));
2539 JsonValue::Object(request)
2540 })
2541 .expect("serialize request"),
2542 );
2543 let response = srv.handle_message(&request).expect("response");
2544 let parsed = parse_json(&response);
2545 let result = parsed.get("result").expect("result");
2546 assert_ne!(
2547 result.get("isError").and_then(JsonValue::as_bool),
2548 Some(true),
2549 "tool returned error: {parsed:?}"
2550 );
2551 parse_json(&tool_result_text(&parsed))
2552 }
2553
2554 #[test]
2555 fn tools_call_explain_connection_maps_supported_schemes() {
2556 let mut srv = make_server();
2557 for (uri, mode, transport, topology) in [
2558 ("memory://", "embedded", "memory", "ephemeral"),
2559 (
2560 "file:///tmp/reddb/data.rdb",
2561 "embedded",
2562 "filesystem",
2563 "persisted",
2564 ),
2565 ("red://db.example:5050", "remote", "redwire", "single"),
2566 ("reds://db.example", "remote", "redwire_tls", "single"),
2567 ("grpc://db.example", "remote", "grpc", "single"),
2568 ("http://db.example", "remote", "http", "single"),
2569 ("https://db.example", "remote", "https", "single"),
2570 ] {
2571 let envelope = call_explain_connection(&mut srv, uri);
2572 assert_eq!(envelope.get("uri").and_then(JsonValue::as_str), Some(uri));
2573 assert_eq!(
2574 envelope.get("mode").and_then(JsonValue::as_str),
2575 Some(mode),
2576 "{uri}: {envelope:?}"
2577 );
2578 assert_eq!(
2579 envelope.get("transport").and_then(JsonValue::as_str),
2580 Some(transport),
2581 "{uri}: {envelope:?}"
2582 );
2583 assert_eq!(
2584 envelope.get("topology").and_then(JsonValue::as_str),
2585 Some(topology),
2586 "{uri}: {envelope:?}"
2587 );
2588 }
2589
2590 let cluster = call_explain_connection(&mut srv, "grpc://primary,replica?route=primary");
2591 assert_eq!(
2592 cluster.get("topology").and_then(JsonValue::as_str),
2593 Some("primary_replica")
2594 );
2595 assert_eq!(
2596 cluster.get("transport").and_then(JsonValue::as_str),
2597 Some("grpc")
2598 );
2599 assert_eq!(
2600 cluster
2601 .get("target")
2602 .and_then(|target| target.get("force_primary"))
2603 .and_then(JsonValue::as_bool),
2604 Some(true)
2605 );
2606 }
2607
2608 struct EnvVarGuard {
2609 name: &'static str,
2610 previous: Option<String>,
2611 }
2612
2613 impl EnvVarGuard {
2614 fn set(name: &'static str, value: &str) -> Self {
2615 let previous = std::env::var(name).ok();
2616 std::env::set_var(name, value);
2617 Self { name, previous }
2618 }
2619 }
2620
2621 impl Drop for EnvVarGuard {
2622 fn drop(&mut self) {
2623 if let Some(value) = self.previous.take() {
2624 std::env::set_var(self.name, value);
2625 } else {
2626 std::env::remove_var(self.name);
2627 }
2628 }
2629 }
2630
2631 struct AskStub {
2632 addr: SocketAddr,
2633 shutdown: Arc<AtomicBool>,
2634 handle: Option<JoinHandle<()>>,
2635 }
2636
2637 impl AskStub {
2638 fn start() -> Self {
2639 let listener = TcpListener::bind("127.0.0.1:0").expect("stub bind");
2640 listener
2641 .set_nonblocking(true)
2642 .expect("nonblocking listener");
2643 let addr = listener.local_addr().expect("local addr");
2644 let shutdown = Arc::new(AtomicBool::new(false));
2645 let server_shutdown = Arc::clone(&shutdown);
2646 let handle = thread::spawn(move || {
2647 while !server_shutdown.load(Ordering::Relaxed) {
2648 match listener.accept() {
2649 Ok((mut stream, _)) => {
2650 let request = read_stub_request(&mut stream);
2651 if request.contains("/embeddings") {
2652 write_json_response(
2653 &mut stream,
2654 r#"{"model":"mock-embedding","data":[{"index":0,"embedding":[1,0,0]}],"usage":{"prompt_tokens":3,"total_tokens":3}}"#,
2655 );
2656 } else {
2657 write_json_response(
2658 &mut stream,
2659 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}}"#,
2660 );
2661 }
2662 }
2663 Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => {
2664 thread::sleep(Duration::from_millis(1));
2665 }
2666 Err(_) => break,
2667 }
2668 }
2669 });
2670
2671 Self {
2672 addr,
2673 shutdown,
2674 handle: Some(handle),
2675 }
2676 }
2677
2678 fn addr(&self) -> SocketAddr {
2679 self.addr
2680 }
2681 }
2682
2683 impl Drop for AskStub {
2684 fn drop(&mut self) {
2685 self.shutdown.store(true, Ordering::Relaxed);
2686 let _ = TcpStream::connect(self.addr);
2687 if let Some(handle) = self.handle.take() {
2688 let _ = handle.join();
2689 }
2690 }
2691 }
2692
2693 fn read_stub_request(stream: &mut TcpStream) -> String {
2694 let _ = stream.set_read_timeout(Some(Duration::from_millis(100)));
2695 let mut buffer = [0_u8; 4096];
2696 let count = stream.read(&mut buffer).unwrap_or(0);
2697 String::from_utf8_lossy(&buffer[..count]).into_owned()
2698 }
2699
2700 fn write_json_response(stream: &mut TcpStream, body: &str) {
2701 let response = format!(
2702 "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
2703 body.len()
2704 );
2705 stream
2706 .write_all(response.as_bytes())
2707 .expect("write stub response");
2708 }
2709}