reddb_server/mcp/

server.rs

//! MCP Server for RedDB.
//!
//! Runs an embedded RedDB runtime and exposes it to AI agents via the
//! Model Context Protocol JSON-RPC transport over stdio.

use crate::application::{
    CreateDocumentInput, CreateEdgeInput, CreateNodeInput, CreateRowInput, CreateVectorInput,
    DeleteEntityInput, EntityUseCases, ExecuteQueryInput, GraphCentralityInput,
    GraphClusteringInput, GraphCommunitiesInput, GraphComponentsInput, GraphCyclesInput,
    GraphShortestPathInput, GraphTraversalInput, GraphUseCases, ScanCollectionInput,
    SearchSimilarInput, SearchTextInput,
};
use crate::auth::store::AuthStore;
use crate::auth::{AuthConfig, Role};
use crate::json::{
    from_str as json_from_str, to_string as json_to_string, Map, Value as JsonValue,
};
use crate::mcp::capability::McpCapability;
use crate::mcp::{protocol, tools};
use crate::presentation::entity_json::created_entity_output_json;
use crate::presentation::entity_json::storage_value_to_json;
use crate::presentation::query_result_json::{runtime_query_json, runtime_stats_json};
use crate::runtime::{
    query_request::ParamValue, RedDBRuntime, RuntimeGraphCentralityAlgorithm,
    RuntimeGraphCommunityAlgorithm, RuntimeGraphComponentsMode, RuntimeGraphDirection,
    RuntimeGraphPathAlgorithm, RuntimeGraphTraversalStrategy,
};
use crate::storage::EntityId;
use reddb_types::Value;

use std::io::{self, BufRead, Write};
use std::sync::Arc;

/// MCP server wrapping an embedded RedDB runtime.
pub struct McpServer {
    runtime: RedDBRuntime,
    auth_store: Arc<AuthStore>,
    initialized: bool,
    /// What this session may do. Read-only unless the operator opted in —
    /// see [`crate::mcp::capability`] for why the tool list, not the
    /// transport, is the boundary that matters here.
    capability: McpCapability,
}

impl McpServer {
    /// Create a new MCP server with the given runtime.
    pub fn new(runtime: RedDBRuntime) -> Self {
        let auth_store = Arc::new(AuthStore::new(AuthConfig {
            enabled: true,
            ..Default::default()
        }));
        auth_store.bootstrap_from_env();
        runtime.set_auth_store(Arc::clone(&auth_store));
        Self {
            runtime,
            auth_store,
            initialized: false,
            capability: McpCapability::default(),
        }
    }

    /// Set the session's capability level. The default is read-only; the
    /// CLI raises it from `--allow-write` / `--allow-admin`.
    pub fn with_capability(mut self, capability: McpCapability) -> Self {
        self.capability = capability;
        self
    }

    /// The session's capability level.
    pub fn capability(&self) -> McpCapability {
        self.capability
    }

    /// Refuse `tool_name` when the session is below the level it needs.
    fn check_tool_capability(&self, tool_name: &str) -> Result<(), String> {
        let required = crate::mcp::capability::tool_capability(tool_name);
        if self.capability.permits(required) {
            Ok(())
        } else {
            Err(crate::mcp::capability::refusal(
                &format!("tool `{tool_name}`"),
                self.capability,
                required,
            ))
        }
    }

    /// Run the MCP server reading from stdin and writing to stdout.
    ///
    /// This blocks until stdin is closed (EOF). Diagnostic messages are
    /// written to stderr so they do not interfere with the protocol.
    pub fn run_stdio(&mut self) {
        let stdin = io::stdin();
        let stdout = io::stdout();
        let mut reader = io::BufReader::new(stdin.lock());
        let mut writer = io::BufWriter::new(stdout.lock());

        tracing::info!(target: "reddb::mcp", "server started, waiting for messages on stdin");

        loop {
            let payload = match protocol::read_payload(&mut reader) {
                Ok(Some(p)) => p,
                Ok(None) => {
                    tracing::info!(target: "reddb::mcp", "stdin closed, shutting down");
                    break;
                }
                Err(e) => {
                    tracing::error!(target: "reddb::mcp", err = %e, "read error");
                    continue;
                }
            };

            let request: JsonValue = match json_from_str(&payload) {
                Ok(v) => v,
                Err(e) => {
                    tracing::warn!(target: "reddb::mcp", err = %e, "invalid JSON");
                    let msg = protocol::build_error_message(None, -32700, "parse error");
                    let _ = protocol::write_message(&mut writer, &msg);
                    continue;
                }
            };

            let response = self.handle_message(&request);
            if let Some(resp) = response {
                if let Err(e) = protocol::write_message(&mut writer, &resp) {
                    tracing::error!(target: "reddb::mcp", err = %e, "write error");
                    break;
                }
            }
        }
    }

    /// Route a JSON-RPC message to the appropriate handler.
    fn handle_message(&mut self, msg: &JsonValue) -> Option<String> {
        let method = msg.get("method").and_then(|v| v.as_str())?;
        let id = msg.get("id");

        match method {
            "initialize" => Some(self.handle_initialize(id)),
            "initialized" | "notifications/initialized" => {
                // Notification -- no response required.
                None
            }
            "tools/list" => Some(self.handle_tools_list(id)),
            "tools/call" => Some(self.handle_tools_call(id, msg.get("params"))),
            "resources/list" => Some(self.handle_resources_list(id)),
            "resources/read" => Some(self.handle_resources_read(id, msg.get("params"))),
            "ping" => {
                let mut result = Map::new();
                result.insert("status".to_string(), JsonValue::String("ok".to_string()));
                Some(protocol::build_result_message(
                    id,
                    JsonValue::Object(result),
                ))
            }
            _ => Some(protocol::build_error_message(
                id,
                -32601,
                &format!("unknown method: {}", method),
            )),
        }
    }

    // ------------------------------------------------------------------
    // MCP lifecycle
    // ------------------------------------------------------------------

    fn handle_initialize(&mut self, id: Option<&JsonValue>) -> String {
        self.initialized = true;

        let mut capabilities = Map::new();
        {
            let mut tools_cap = Map::new();
            tools_cap.insert("listChanged".to_string(), JsonValue::Bool(false));
            capabilities.insert("tools".to_string(), JsonValue::Object(tools_cap));
        }
        {
            // Read-only knowledge resources (ADR 0061). No subscriptions and
            // the set is static, so both change-notification flags are false.
            let mut resources_cap = Map::new();
            resources_cap.insert("subscribe".to_string(), JsonValue::Bool(false));
            resources_cap.insert("listChanged".to_string(), JsonValue::Bool(false));
            capabilities.insert("resources".to_string(), JsonValue::Object(resources_cap));
        }

        let mut server_info = Map::new();
        server_info.insert(
            "name".to_string(),
            JsonValue::String("reddb-mcp".to_string()),
        );
        server_info.insert(
            "version".to_string(),
            JsonValue::String(env!("CARGO_PKG_VERSION").to_string()),
        );

        let mut result = Map::new();
        result.insert(
            "protocolVersion".to_string(),
            JsonValue::String("2024-11-05".to_string()),
        );
        result.insert("capabilities".to_string(), JsonValue::Object(capabilities));
        result.insert("serverInfo".to_string(), JsonValue::Object(server_info));

        protocol::build_result_message(id, JsonValue::Object(result))
    }

    // ------------------------------------------------------------------
    // tools/list
    // ------------------------------------------------------------------

    fn handle_tools_list(&self, id: Option<&JsonValue>) -> String {
        let defs = tools::all_tools();
        let mut tools_json: Vec<JsonValue> = defs
            .into_iter()
            .map(|def| {
                let mut obj = Map::new();
                obj.insert("name".to_string(), JsonValue::String(def.name.to_string()));
                obj.insert(
                    "description".to_string(),
                    JsonValue::String(def.descriptor_description()),
                );
                obj.insert("inputSchema".to_string(), def.input_schema);
                JsonValue::Object(obj)
            })
            .collect();
        tools_json.push(tools::ask_descriptor());

        let mut result = Map::new();
        result.insert("tools".to_string(), JsonValue::Array(tools_json));
        protocol::build_result_message(id, JsonValue::Object(result))
    }

    // ------------------------------------------------------------------
    // resources/list + resources/read (knowledge surface, ADR 0061)
    // ------------------------------------------------------------------

    fn handle_resources_list(&self, id: Option<&JsonValue>) -> String {
        let resources: Vec<JsonValue> = tools::knowledge_resources()
            .iter()
            .map(|res| {
                let mut obj = Map::new();
                obj.insert("uri".to_string(), JsonValue::String(res.uri.to_string()));
                obj.insert("name".to_string(), JsonValue::String(res.title.to_string()));
                obj.insert(
                    "description".to_string(),
                    JsonValue::String(res.description.to_string()),
                );
                obj.insert(
                    "mimeType".to_string(),
                    JsonValue::String(res.mime_type.to_string()),
                );
                JsonValue::Object(obj)
            })
            .collect();

        let mut result = Map::new();
        result.insert("resources".to_string(), JsonValue::Array(resources));
        protocol::build_result_message(id, JsonValue::Object(result))
    }

    fn handle_resources_read(&self, id: Option<&JsonValue>, params: Option<&JsonValue>) -> String {
        let uri = params.and_then(|p| p.get("uri")).and_then(|v| v.as_str());
        let uri = match uri {
            Some(u) => u,
            None => return protocol::build_error_message(id, -32602, "missing resource uri"),
        };

        let resources = tools::knowledge_resources();
        let resource = match resources.iter().find(|res| res.uri == uri) {
            Some(res) => res,
            None => {
                return protocol::build_error_message(
                    id,
                    -32602,
                    &format!("unknown resource: {uri}"),
                );
            }
        };

        let mut contents = Map::new();
        contents.insert(
            "uri".to_string(),
            JsonValue::String(resource.uri.to_string()),
        );
        contents.insert(
            "mimeType".to_string(),
            JsonValue::String(resource.mime_type.to_string()),
        );
        contents.insert("text".to_string(), JsonValue::String((resource.body)()));

        let mut result = Map::new();
        result.insert(
            "contents".to_string(),
            JsonValue::Array(vec![JsonValue::Object(contents)]),
        );
        protocol::build_result_message(id, JsonValue::Object(result))
    }

    // ------------------------------------------------------------------
    // tools/call dispatcher
    // ------------------------------------------------------------------

    fn handle_tools_call(&self, id: Option<&JsonValue>, params: Option<&JsonValue>) -> String {
        let name = params.and_then(|p| p.get("name")).and_then(|v| v.as_str());
        let name = match name {
            Some(n) => n,
            None => {
                return protocol::build_error_message(id, -32602, "missing tool name");
            }
        };

        let empty = JsonValue::Object(Map::new());
        let args = params.and_then(|p| p.get("arguments")).unwrap_or(&empty);

        let result = tools::authorize_tool(name, tools::McpIdentityPosture::EmbeddedImplicitAdmin)
            .and_then(|_| self.check_tool_capability(name))
            .and_then(|_| match name {
                "reddb_query" => self.tool_query(args),
                "reddb_collections" => self.tool_collections(),
                "reddb_insert_row" => self.tool_insert_row(args),
                "reddb_insert_node" => self.tool_insert_node(args),
                "reddb_insert_edge" => self.tool_insert_edge(args),
                "reddb_insert_vector" => self.tool_insert_vector(args),
                "reddb_insert_document" => self.tool_insert_document(args),
                "reddb_kv_get" => self.tool_kv_get(args),
                "reddb_kv_set" => self.tool_kv_set(args),
                "reddb_kv_invalidate_tags" => self.tool_kv_invalidate_tags(args),
                "reddb_config_get" => self.tool_config_get(args),
                "reddb_config_put" => self.tool_config_put(args),
                "reddb_config_resolve" => self.tool_config_resolve(args),
                "reddb_vault_get" => self.tool_vault_get(args),
                "reddb_vault_put" => self.tool_vault_put(args),
                "reddb_vault_reveal" | "reddb_vault_unseal" => self.tool_vault_unseal(args),
                "reddb_delete" => self.tool_delete(args),
                "reddb_search_vector" => self.tool_search_vector(args),
                "reddb_search_text" => self.tool_search_text(args),
                "reddb_health" => self.tool_health(),
                "reddb_graph_traverse" => self.tool_graph_traverse(args),
                "reddb_graph_shortest_path" => self.tool_graph_shortest_path(args),
                "reddb_update" => self.tool_update(args),
                "reddb_scan" => self.tool_scan(args),
                "reddb_graph_centrality" => self.tool_graph_centrality(args),
                "reddb_graph_community" => self.tool_graph_community(args),
                "reddb_graph_components" => self.tool_graph_components(args),
                "reddb_graph_cycles" => self.tool_graph_cycles(args),
                "reddb_graph_clustering" => self.tool_graph_clustering(args),
                "reddb_create_collection" => self.tool_create_collection(args),
                "reddb_drop_collection" => self.tool_drop_collection(args),
                "reddb_rql_validate" => self.tool_rql_validate(args),
                "reddb_rql_explain" => self.tool_rql_explain(args),
                "reddb_type_of" => self.tool_type_of(args),
                "reddb_explain_connection" => self.tool_explain_connection(args),
                "reddb_auth_bootstrap" => self.tool_auth_bootstrap(args),
                "reddb_auth_create_user" => self.tool_auth_create_user(args),
                "reddb_auth_login" => self.tool_auth_login(args),
                "reddb_auth_create_api_key" => self.tool_auth_create_api_key(args),
                "reddb_auth_list_users" => self.tool_auth_list_users(),
                crate::runtime::ai::mcp_ask_tool::TOOL_NAME => self.tool_ask(args),
                _ => Err(format!("unknown tool: {name}")),
            });

        match result {
            Ok(text) => {
                let mut content = Map::new();
                content.insert("type".to_string(), JsonValue::String("text".to_string()));
                content.insert("text".to_string(), JsonValue::String(text));

                let mut result_obj = Map::new();
                result_obj.insert(
                    "content".to_string(),
                    JsonValue::Array(vec![JsonValue::Object(content)]),
                );
                protocol::build_result_message(id, JsonValue::Object(result_obj))
            }
            Err(err) => {
                let mut content = Map::new();
                content.insert("type".to_string(), JsonValue::String("text".to_string()));
                content.insert("text".to_string(), JsonValue::String(err.clone()));

                let mut result_obj = Map::new();
                result_obj.insert(
                    "content".to_string(),
                    JsonValue::Array(vec![JsonValue::Object(content)]),
                );
                result_obj.insert("isError".to_string(), JsonValue::Bool(true));
                protocol::build_result_message(id, JsonValue::Object(result_obj))
            }
        }
    }

    // ------------------------------------------------------------------
    // Tool implementations
    // ------------------------------------------------------------------

    fn tool_query(&self, args: &JsonValue) -> Result<String, String> {
        let sql = args
            .get("sql")
            .and_then(|v| v.as_str())
            .ok_or("missing required field 'sql'")?;

        // `reddb_query` is a read-only *tool* that accepts arbitrary SQL, so
        // the statement decides what it needs. Without this a read-only
        // session could send `DROP COLLECTION` through the one tool it was
        // allowed to call.
        let required = crate::mcp::capability::statement_capability(sql);
        if !self.capability.permits(required) {
            return Err(crate::mcp::capability::refusal(
                "this statement",
                self.capability,
                required,
            ));
        }

        // Optional positional `$N` bind parameters. MCP, stdio, and HTTP
        // all decode through the Request module's transport-neutral codec.
        if let Some(raw_params) = args.get("params") {
            let arr = raw_params
                .as_array()
                .ok_or_else(|| "'params' must be an array".to_string())?;
            let binds: Vec<Value> = arr
                .iter()
                .map(ParamValue::decode_json)
                .collect::<Result<Vec<_>, _>>()?
                .into_iter()
                .map(Value::from)
                .collect();

            // Bind inside the runtime's statement frame (#2183) so a
            // parameterized MCP query keeps the same snapshot isolation and
            // `AS OF` resolution as textual SQL.
            let result = self
                .runtime
                .execute_query_with_params(sql, &binds)
                .map_err(|e| format!("{}", e))?;
            let json = runtime_query_json(&result, &None, &None);
            return json_to_string(&json).map_err(|e| format!("serialization error: {}", e));
        }

        let uc = &self.runtime;
        let result = uc.execute_query(sql).map_err(|e| format!("{}", e))?;

        let json = runtime_query_json(&result, &None, &None);
        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_ask(&self, args: &JsonValue) -> Result<String, String> {
        let invocation =
            crate::runtime::ai::mcp_ask_tool::parse(args).map_err(format_mcp_ask_parse_error)?;
        let ask = reddb_rql::ast::AskQuery {
            explain: false,
            question: invocation.question,
            question_param: None,
            provider: invocation.using,
            model: invocation.model,
            depth: invocation.depth.map(|v| v as usize),
            limit: invocation.limit.map(|v| v as usize),
            min_score: invocation.min_score.map(|v| v as f32),
            collection: None,
            temperature: invocation.temperature.map(|v| v as f32),
            seed: invocation.seed,
            strict: invocation.strict.unwrap_or(true),
            stream: false,
            cache: if matches!(invocation.nocache, Some(true)) {
                reddb_rql::ast::AskCacheClause::NoCache
            } else if let Some(ttl) = invocation.cache_ttl {
                reddb_rql::ast::AskCacheClause::CacheTtl(ttl)
            } else {
                reddb_rql::ast::AskCacheClause::Default
            },
            plan_only: false,
            steps: None,
        };
        let result = self
            .runtime
            .execute_ask("ASK <mcp>", &ask)
            .map_err(|e| format!("{}", e))?;
        let json = crate::presentation::query_result::tagged_summary(&result);
        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_collections(&self) -> Result<String, String> {
        let uc = &self.runtime;
        let collections = uc.db().collections();
        let json = JsonValue::Array(collections.into_iter().map(JsonValue::String).collect());
        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_insert_row(&self, args: &JsonValue) -> Result<String, String> {
        let collection = args
            .get("collection")
            .and_then(|v| v.as_str())
            .ok_or("missing required field 'collection'")?;
        let data = args
            .get("data")
            .and_then(|v| v.as_object())
            .ok_or("missing required field 'data' (must be an object)")?;

        let mut fields = Vec::new();
        for (key, value) in data {
            let sv = crate::application::entity::json_to_storage_value(value)
                .map_err(|e| format!("{}", e))?;
            fields.push((key.clone(), sv));
        }

        let metadata = parse_metadata_arg(args)?;

        let uc = EntityUseCases::new(&self.runtime);
        let output = uc
            .create_row(CreateRowInput {
                collection: collection.to_string(),
                fields,
                metadata,
                node_links: vec![],
                vector_links: vec![],
            })
            .map_err(|e| format!("{}", e))?;

        let json = created_entity_output_json(&output);
        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_insert_node(&self, args: &JsonValue) -> Result<String, String> {
        let collection = args
            .get("collection")
            .and_then(|v| v.as_str())
            .ok_or("missing required field 'collection'")?;
        let label = args
            .get("label")
            .and_then(|v| v.as_str())
            .ok_or("missing required field 'label'")?;
        let node_type = args
            .get("node_type")
            .and_then(|v| v.as_str())
            .map(String::from);

        let mut properties = Vec::new();
        if let Some(props) = args.get("properties").and_then(|v| v.as_object()) {
            for (key, value) in props {
                let sv = crate::application::entity::json_to_storage_value(value)
                    .map_err(|e| format!("{}", e))?;
                properties.push((key.clone(), sv));
            }
        }

        let metadata = parse_metadata_arg(args)?;

        let uc = EntityUseCases::new(&self.runtime);
        let output = uc
            .create_node(CreateNodeInput {
                collection: collection.to_string(),
                label: label.to_string(),
                node_type,
                properties,
                metadata,
                embeddings: vec![],
                table_links: vec![],
                node_links: vec![],
            })
            .map_err(|e| format!("{}", e))?;

        let json = created_entity_output_json(&output);
        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_insert_edge(&self, args: &JsonValue) -> Result<String, String> {
        let collection = args
            .get("collection")
            .and_then(|v| v.as_str())
            .ok_or("missing required field 'collection'")?;
        let label = args
            .get("label")
            .and_then(|v| v.as_str())
            .ok_or("missing required field 'label'")?;
        let from_id = args
            .get("from")
            .and_then(|v| v.as_u64())
            .ok_or("missing required field 'from' (integer)")?;
        let to_id = args
            .get("to")
            .and_then(|v| v.as_u64())
            .ok_or("missing required field 'to' (integer)")?;
        let weight = args
            .get("weight")
            .and_then(|v| v.as_f64())
            .map(|w| w as f32);

        let mut properties = Vec::new();
        if let Some(props) = args.get("properties").and_then(|v| v.as_object()) {
            for (key, value) in props {
                let sv = crate::application::entity::json_to_storage_value(value)
                    .map_err(|e| format!("{}", e))?;
                properties.push((key.clone(), sv));
            }
        }

        let metadata = parse_metadata_arg(args)?;

        let uc = EntityUseCases::new(&self.runtime);
        let output = uc
            .create_edge(CreateEdgeInput {
                collection: collection.to_string(),
                label: label.to_string(),
                from: EntityId::new(from_id),
                to: EntityId::new(to_id),
                weight,
                properties,
                metadata,
            })
            .map_err(|e| format!("{}", e))?;

        let json = created_entity_output_json(&output);
        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_insert_vector(&self, args: &JsonValue) -> Result<String, String> {
        let collection = args
            .get("collection")
            .and_then(|v| v.as_str())
            .ok_or("missing required field 'collection'")?;
        let dense_arr = args
            .get("dense")
            .and_then(|v| v.as_array())
            .ok_or("missing required field 'dense' (array of numbers)")?;

        let mut dense = Vec::with_capacity(dense_arr.len());
        for v in dense_arr {
            dense.push(
                v.as_f64()
                    .ok_or("'dense' array must contain only numbers")? as f32,
            );
        }
        if dense.is_empty() {
            return Err("'dense' vector cannot be empty".to_string());
        }

        let content = args
            .get("content")
            .and_then(|v| v.as_str())
            .map(String::from);
        let metadata = parse_metadata_arg(args)?;

        let uc = EntityUseCases::new(&self.runtime);
        let output = uc
            .create_vector(CreateVectorInput {
                collection: collection.to_string(),
                dense,
                content,
                metadata,
                link_row: None,
                link_node: None,
            })
            .map_err(|e| format!("{}", e))?;

        let json = created_entity_output_json(&output);
        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_insert_document(&self, args: &JsonValue) -> Result<String, String> {
        let collection = args
            .get("collection")
            .and_then(|v| v.as_str())
            .ok_or("missing required field 'collection'")?;
        let body = args.get("body").ok_or("missing required field 'body'")?;

        let metadata = parse_metadata_arg(args)?;

        let uc = EntityUseCases::new(&self.runtime);
        let output = uc
            .create_document(CreateDocumentInput {
                collection: collection.to_string(),
                body: body.clone(),
                metadata,
                node_links: vec![],
                vector_links: vec![],
            })
            .map_err(|e| format!("{}", e))?;

        let json = created_entity_output_json(&output);
        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_kv_get(&self, args: &JsonValue) -> Result<String, String> {
        let collection = args
            .get("collection")
            .and_then(|v| v.as_str())
            .ok_or("missing required field 'collection'")?;
        let key = args
            .get("key")
            .and_then(|v| v.as_str())
            .ok_or("missing required field 'key'")?;

        let uc = EntityUseCases::new(&self.runtime);
        match uc.get_kv(collection, key).map_err(|e| format!("{}", e))? {
            Some((value, entity_id)) => {
                let mut obj = Map::new();
                obj.insert("found".to_string(), JsonValue::Bool(true));
                obj.insert("key".to_string(), JsonValue::String(key.to_string()));
                obj.insert("value".to_string(), storage_value_to_json(&value));
                obj.insert("rid".to_string(), JsonValue::Number(entity_id.raw() as f64));
                obj.insert("kind".to_string(), JsonValue::String("kv".to_string()));
                json_to_string(&JsonValue::Object(obj))
                    .map_err(|e| format!("serialization error: {}", e))
            }
            None => {
                let mut obj = Map::new();
                obj.insert("found".to_string(), JsonValue::Bool(false));
                obj.insert("key".to_string(), JsonValue::String(key.to_string()));
                json_to_string(&JsonValue::Object(obj))
                    .map_err(|e| format!("serialization error: {}", e))
            }
        }
    }

    fn tool_kv_set(&self, args: &JsonValue) -> Result<String, String> {
        let collection = args
            .get("collection")
            .and_then(|v| v.as_str())
            .ok_or("missing required field 'collection'")?;
        let key = args
            .get("key")
            .and_then(|v| v.as_str())
            .ok_or("missing required field 'key'")?;
        let value_arg = args.get("value").ok_or("missing required field 'value'")?;

        let sv = crate::application::entity::json_to_storage_value(value_arg)
            .map_err(|e| format!("{}", e))?;

        let metadata = parse_metadata_arg(args)?;

        let tags = parse_string_array_arg(args, "tags")?;
        let ops = crate::runtime::impl_kv::KvAtomicOps::new(&self.runtime);
        let (_, id) = ops
            .set_with_tags_and_metadata(collection, key, sv, None, &tags, false, metadata)
            .map_err(|e| format!("{}", e))?;

        let mut obj = Map::new();
        obj.insert("ok".to_string(), JsonValue::Bool(true));
        obj.insert("rid".to_string(), JsonValue::Number(id.raw() as f64));
        obj.insert("kind".to_string(), JsonValue::String("kv".to_string()));
        obj.insert(
            "tags".to_string(),
            JsonValue::Array(tags.into_iter().map(JsonValue::String).collect()),
        );
        json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_kv_invalidate_tags(&self, args: &JsonValue) -> Result<String, String> {
        let collection = args
            .get("collection")
            .and_then(|v| v.as_str())
            .ok_or("missing required field 'collection'")?;
        let tags = parse_string_array_arg(args, "tags")?;
        if tags.is_empty() {
            return Err("missing required field 'tags'".to_string());
        }
        let ops = crate::runtime::impl_kv::KvAtomicOps::new(&self.runtime);
        let count = ops
            .invalidate_tags(collection, &tags)
            .map_err(|e| format!("{}", e))?;

        let mut obj = Map::new();
        obj.insert("ok".to_string(), JsonValue::Bool(true));
        obj.insert("invalidated".to_string(), JsonValue::Number(count as f64));
        obj.insert(
            "tags".to_string(),
            JsonValue::Array(tags.into_iter().map(JsonValue::String).collect()),
        );
        json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_config_get(&self, args: &JsonValue) -> Result<String, String> {
        let collection = mcp_keyed_ident(get_str_field(args, "collection")?)?;
        let key = mcp_keyed_ident(get_str_field(args, "key")?)?;
        self.tool_keyed_query(format!("GET CONFIG {collection} {key}"))
    }

    fn tool_config_put(&self, args: &JsonValue) -> Result<String, String> {
        reject_mcp_volatile_options(args, "CONFIG")?;
        let collection = mcp_keyed_ident(get_str_field(args, "collection")?)?;
        let key = mcp_keyed_ident(get_str_field(args, "key")?)?;
        let tags = parse_string_array_arg(args, "tags")?;
        let literal = if let Some(secret_ref) = args.get("secret_ref") {
            let object = secret_ref
                .as_object()
                .ok_or("field 'secret_ref' must be an object")?;
            let ref_collection = object
                .get("collection")
                .and_then(|v| v.as_str())
                .ok_or_else(|| "secret_ref.collection is required".to_string())
                .and_then(mcp_keyed_ident)?;
            let ref_key = object
                .get("key")
                .and_then(|v| v.as_str())
                .ok_or_else(|| "secret_ref.key is required".to_string())
                .and_then(mcp_keyed_ident)?;
            format!("SECRET_REF(vault, {ref_collection}.{ref_key})")
        } else {
            mcp_value_literal(args.get("value").ok_or("missing required field 'value'")?)?
        };
        let mut sql = format!("PUT CONFIG {collection} {key} = {literal}");
        append_mcp_tags_clause(&mut sql, &tags);
        self.tool_keyed_query(sql)
    }

    fn tool_config_resolve(&self, args: &JsonValue) -> Result<String, String> {
        let collection = mcp_keyed_ident(get_str_field(args, "collection")?)?;
        let key = mcp_keyed_ident(get_str_field(args, "key")?)?;
        self.tool_keyed_query(format!("RESOLVE CONFIG {collection} {key}"))
    }

    fn tool_vault_get(&self, args: &JsonValue) -> Result<String, String> {
        let collection = mcp_keyed_ident(get_str_field(args, "collection")?)?;
        let key = mcp_keyed_ident(get_str_field(args, "key")?)?;
        self.tool_keyed_query(format!("VAULT GET {collection}.{key}"))
    }

    fn tool_vault_put(&self, args: &JsonValue) -> Result<String, String> {
        reject_mcp_volatile_options(args, "VAULT")?;
        let collection = mcp_keyed_ident(get_str_field(args, "collection")?)?;
        let key = mcp_keyed_ident(get_str_field(args, "key")?)?;
        let value = mcp_value_literal(args.get("value").ok_or("missing required field 'value'")?)?;
        let tags = parse_string_array_arg(args, "tags")?;
        let mut sql = format!("VAULT PUT {collection}.{key} = {value}");
        append_mcp_tags_clause(&mut sql, &tags);
        self.tool_keyed_query(sql)
    }

    fn tool_vault_unseal(&self, args: &JsonValue) -> Result<String, String> {
        let collection = mcp_keyed_ident(get_str_field(args, "collection")?)?;
        let key = mcp_keyed_ident(get_str_field(args, "key")?)?;
        self.tool_keyed_query(format!("VAULT REVEAL {collection}.{key}"))
    }

    fn tool_keyed_query(&self, sql: String) -> Result<String, String> {
        let uc = &self.runtime;
        let result = uc.execute_query(&(sql)).map_err(|e| format!("{}", e))?;
        let json = runtime_query_json(&result, &None, &None);
        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_delete(&self, args: &JsonValue) -> Result<String, String> {
        let collection = args
            .get("collection")
            .and_then(|v| v.as_str())
            .ok_or("missing required field 'collection'")?;
        let id = args
            .get("id")
            .and_then(|v| v.as_u64())
            .ok_or("missing required field 'id' (integer)")?;

        let uc = EntityUseCases::new(&self.runtime);
        let output = uc
            .delete(DeleteEntityInput {
                collection: collection.to_string(),
                id: EntityId::new(id),
            })
            .map_err(|e| format!("{}", e))?;

        let mut obj = Map::new();
        obj.insert("deleted".to_string(), JsonValue::Bool(output.deleted));
        obj.insert("id".to_string(), JsonValue::Number(output.id.raw() as f64));
        json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_search_vector(&self, args: &JsonValue) -> Result<String, String> {
        let collection = args
            .get("collection")
            .and_then(|v| v.as_str())
            .ok_or("missing required field 'collection'")?;
        let vector_arr = args
            .get("vector")
            .and_then(|v| v.as_array())
            .ok_or("missing required field 'vector' (array of numbers)")?;

        let mut vector = Vec::with_capacity(vector_arr.len());
        for v in vector_arr {
            vector.push(
                v.as_f64()
                    .ok_or("'vector' array must contain only numbers")? as f32,
            );
        }
        let k = args
            .get("k")
            .and_then(|v| v.as_u64())
            .map(|v| v as usize)
            .unwrap_or(10);
        let min_score = args
            .get("min_score")
            .and_then(|v| v.as_f64())
            .map(|v| v as f32)
            .unwrap_or(0.0);

        let uc = &self.runtime;
        let results = uc
            .search_similar_input(SearchSimilarInput {
                collection: collection.to_string(),
                vector,
                k,
                min_score,
                text: None,
                provider: None,
            })
            .map_err(|e| format!("{}", e))?;

        let items: Vec<JsonValue> = results
            .iter()
            .map(|r| {
                let mut obj = Map::new();
                obj.insert(
                    "rid".to_string(),
                    JsonValue::Number(r.entity_id.raw() as f64),
                );
                obj.insert("kind".to_string(), JsonValue::String("vector".to_string()));
                obj.insert("score".to_string(), JsonValue::Number(r.score as f64));
                obj.insert("distance".to_string(), JsonValue::Number(r.distance as f64));
                JsonValue::Object(obj)
            })
            .collect();

        let mut obj = Map::new();
        obj.insert("count".to_string(), JsonValue::Number(items.len() as f64));
        obj.insert("results".to_string(), JsonValue::Array(items));
        json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_search_text(&self, args: &JsonValue) -> Result<String, String> {
        let query = args
            .get("query")
            .and_then(|v| v.as_str())
            .ok_or("missing required field 'query'")?;

        let collections = args
            .get("collections")
            .and_then(|v| v.as_array())
            .map(|arr| {
                arr.iter()
                    .filter_map(|v| v.as_str().map(String::from))
                    .collect::<Vec<_>>()
            });
        let limit = args
            .get("limit")
            .and_then(|v| v.as_u64())
            .map(|v| v as usize);
        let fuzzy = args.get("fuzzy").and_then(|v| v.as_bool()).unwrap_or(false);

        let uc = &self.runtime;
        let result = uc
            .search_text_input(SearchTextInput {
                query: query.to_string(),
                collections,
                entity_types: None,
                capabilities: None,
                fields: None,
                limit,
                fuzzy,
            })
            .map_err(|e| format!("{}", e))?;

        let items: Vec<JsonValue> = result
            .matches
            .iter()
            .map(|m| {
                let mut obj = Map::new();
                obj.insert(
                    "rid".to_string(),
                    JsonValue::Number(m.entity.id.raw() as f64),
                );
                obj.insert(
                    "kind".to_string(),
                    JsonValue::String(m.entity.kind.storage_type().to_string()),
                );
                obj.insert("score".to_string(), JsonValue::Number(m.score as f64));
                JsonValue::Object(obj)
            })
            .collect();

        let mut obj = Map::new();
        obj.insert("count".to_string(), JsonValue::Number(items.len() as f64));
        obj.insert("results".to_string(), JsonValue::Array(items));
        json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_health(&self) -> Result<String, String> {
        let uc = &self.runtime;
        let stats = uc.stats();
        let json = runtime_stats_json(&stats);
        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_graph_traverse(&self, args: &JsonValue) -> Result<String, String> {
        let source = args
            .get("source")
            .and_then(|v| v.as_str())
            .ok_or("missing required field 'source'")?;
        let direction = parse_direction(args.get("direction").and_then(|v| v.as_str()));
        let max_depth = args
            .get("max_depth")
            .and_then(|v| v.as_u64())
            .map(|v| v as usize)
            .unwrap_or(3);
        let strategy = match args.get("strategy").and_then(|v| v.as_str()) {
            Some("dfs") => RuntimeGraphTraversalStrategy::Dfs,
            _ => RuntimeGraphTraversalStrategy::Bfs,
        };

        let uc = GraphUseCases::new(&self.runtime);
        let result = uc
            .traverse(GraphTraversalInput {
                source: source.to_string(),
                direction,
                max_depth,
                strategy,
                edge_labels: None,
                projection: None,
            })
            .map_err(|e| format!("{}", e))?;

        let visits: Vec<JsonValue> = result
            .visits
            .iter()
            .map(|v| {
                let mut obj = Map::new();
                obj.insert("depth".to_string(), JsonValue::Number(v.depth as f64));
                obj.insert("node_id".to_string(), JsonValue::String(v.node.id.clone()));
                obj.insert("label".to_string(), JsonValue::String(v.node.label.clone()));
                obj.insert(
                    "node_type".to_string(),
                    JsonValue::String(v.node.node_type.clone()),
                );
                JsonValue::Object(obj)
            })
            .collect();

        let edges: Vec<JsonValue> = result
            .edges
            .iter()
            .map(|e| {
                let mut obj = Map::new();
                obj.insert("source".to_string(), JsonValue::String(e.source.clone()));
                obj.insert("target".to_string(), JsonValue::String(e.target.clone()));
                obj.insert(
                    "edge_type".to_string(),
                    JsonValue::String(e.edge_type.clone()),
                );
                obj.insert("weight".to_string(), JsonValue::Number(e.weight as f64));
                JsonValue::Object(obj)
            })
            .collect();

        let mut obj = Map::new();
        obj.insert(
            "source".to_string(),
            JsonValue::String(result.source.clone()),
        );
        obj.insert(
            "visit_count".to_string(),
            JsonValue::Number(visits.len() as f64),
        );
        obj.insert("visits".to_string(), JsonValue::Array(visits));
        obj.insert("edges".to_string(), JsonValue::Array(edges));
        json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_graph_shortest_path(&self, args: &JsonValue) -> Result<String, String> {
        let source = args
            .get("source")
            .and_then(|v| v.as_str())
            .ok_or("missing required field 'source'")?;
        let target = args
            .get("target")
            .and_then(|v| v.as_str())
            .ok_or("missing required field 'target'")?;
        let direction = parse_direction(args.get("direction").and_then(|v| v.as_str()));
        let algorithm = match args.get("algorithm").and_then(|v| v.as_str()) {
            Some("astar") | Some("a*") => RuntimeGraphPathAlgorithm::AStar,
            Some("bellman_ford") | Some("bellmanford") => RuntimeGraphPathAlgorithm::BellmanFord,
            Some("dijkstra") => RuntimeGraphPathAlgorithm::Dijkstra,
            _ => RuntimeGraphPathAlgorithm::Bfs,
        };

        let uc = GraphUseCases::new(&self.runtime);
        let result = uc
            .shortest_path(GraphShortestPathInput {
                source: source.to_string(),
                target: target.to_string(),
                direction,
                algorithm,
                edge_labels: None,
                projection: None,
            })
            .map_err(|e| format!("{}", e))?;

        let mut obj = Map::new();
        obj.insert(
            "source".to_string(),
            JsonValue::String(result.source.clone()),
        );
        obj.insert(
            "target".to_string(),
            JsonValue::String(result.target.clone()),
        );
        obj.insert(
            "nodes_visited".to_string(),
            JsonValue::Number(result.nodes_visited as f64),
        );

        match &result.path {
            Some(path) => {
                obj.insert("found".to_string(), JsonValue::Bool(true));
                obj.insert(
                    "hop_count".to_string(),
                    JsonValue::Number(path.hop_count as f64),
                );
                obj.insert(
                    "total_weight".to_string(),
                    JsonValue::Number(path.total_weight),
                );
                let nodes_json: Vec<JsonValue> = path
                    .nodes
                    .iter()
                    .map(|n| {
                        let mut nobj = Map::new();
                        nobj.insert("id".to_string(), JsonValue::String(n.id.clone()));
                        nobj.insert("label".to_string(), JsonValue::String(n.label.clone()));
                        JsonValue::Object(nobj)
                    })
                    .collect();
                obj.insert("nodes".to_string(), JsonValue::Array(nodes_json));
            }
            None => {
                obj.insert("found".to_string(), JsonValue::Bool(false));
            }
        }

        json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_update(&self, args: &JsonValue) -> Result<String, String> {
        let collection = get_str_field(args, "collection")?;
        let set_obj = args.get("set").ok_or("missing 'set'")?;
        let where_clause = args
            .get("where_filter")
            .and_then(|v| v.as_str())
            .unwrap_or("");

        // Build UPDATE SQL and execute via runtime. The collection and the
        // column names are spliced as identifiers, so they must be bare
        // identifiers; the values go through the SQL literal renderer, which
        // escapes quotes and backslashes the way the lexer reads them.
        use crate::server::handlers_query::is_safe_sql_identifier;
        if !is_safe_sql_identifier(collection) {
            return Err(format!(
                "invalid collection name '{collection}': expected [A-Za-z_][A-Za-z0-9_]*"
            ));
        }
        let mut sql = format!("UPDATE {} SET ", collection);
        if let Some(obj) = set_obj.as_object() {
            let assignments: Vec<String> = obj
                .iter()
                .map(|(k, v)| {
                    if !is_safe_sql_identifier(k) {
                        return Err(format!(
                            "invalid column name '{k}': expected [A-Za-z_][A-Za-z0-9_]*"
                        ));
                    }
                    Ok(format!(
                        "{} = {}",
                        k,
                        crate::rpc_stdio::value_to_sql_literal(v)
                    ))
                })
                .collect::<Result<_, String>>()?;
            sql.push_str(&assignments.join(", "));
        } else {
            return Err("'set' must be a JSON object".to_string());
        }
        if !where_clause.is_empty() {
            sql.push_str(&format!(" WHERE {}", where_clause));
        }

        let uc = &self.runtime;
        let result = uc.execute_query(&(sql)).map_err(|e| format!("{}", e))?;

        let mut resp = Map::new();
        resp.insert("ok".into(), JsonValue::Bool(true));
        resp.insert(
            "affected_rows".into(),
            JsonValue::Number(result.affected_rows as f64),
        );
        json_to_string(&JsonValue::Object(resp)).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_scan(&self, args: &JsonValue) -> Result<String, String> {
        let collection = get_str_field(args, "collection")?;
        let limit = args
            .get("limit")
            .and_then(|v| v.as_u64())
            .map(|v| v as usize)
            .unwrap_or(10);
        let offset = args
            .get("offset")
            .and_then(|v| v.as_u64())
            .map(|v| v as usize)
            .unwrap_or(0);

        let uc = &self.runtime;
        let page = uc
            .scan(ScanCollectionInput {
                collection: collection.to_string(),
                offset,
                limit,
            })
            .map_err(|e| format!("{}", e))?;

        let json = crate::presentation::entity_json::scan_page_json(&page);
        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_graph_centrality(&self, args: &JsonValue) -> Result<String, String> {
        let algorithm_str = get_str_field(args, "algorithm")?;
        let algo = match algorithm_str {
            "degree" => RuntimeGraphCentralityAlgorithm::Degree,
            "closeness" => RuntimeGraphCentralityAlgorithm::Closeness,
            "betweenness" => RuntimeGraphCentralityAlgorithm::Betweenness,
            "eigenvector" => RuntimeGraphCentralityAlgorithm::Eigenvector,
            "pagerank" => RuntimeGraphCentralityAlgorithm::PageRank,
            _ => return Err(format!("unknown algorithm: {algorithm_str}")),
        };

        let uc = GraphUseCases::new(&self.runtime);
        let result = uc
            .centrality(GraphCentralityInput {
                algorithm: algo,
                top_k: 100,
                normalize: true,
                max_iterations: None,
                epsilon: None,
                alpha: None,
                projection: None,
            })
            .map_err(|e| format!("{}", e))?;

        let json = crate::presentation::graph_json::graph_centrality_json(&result);
        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_graph_community(&self, args: &JsonValue) -> Result<String, String> {
        let algorithm_str = get_str_field(args, "algorithm")?;
        let algo = match algorithm_str {
            "label_propagation" => RuntimeGraphCommunityAlgorithm::LabelPropagation,
            "louvain" => RuntimeGraphCommunityAlgorithm::Louvain,
            _ => return Err(format!("unknown algorithm: {algorithm_str}")),
        };
        let max_iterations = args
            .get("max_iterations")
            .and_then(|v| v.as_u64())
            .map(|v| v as usize);

        let uc = GraphUseCases::new(&self.runtime);
        let result = uc
            .communities(GraphCommunitiesInput {
                algorithm: algo,
                min_size: 1,
                max_iterations,
                resolution: None,
                projection: None,
            })
            .map_err(|e| format!("{}", e))?;

        let json = crate::presentation::graph_json::graph_community_json(&result);
        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_graph_components(&self, args: &JsonValue) -> Result<String, String> {
        let mode = match args.get("mode").and_then(|v| v.as_str()) {
            Some("strongly_connected") => RuntimeGraphComponentsMode::Strong,
            _ => RuntimeGraphComponentsMode::Weak,
        };

        let uc = GraphUseCases::new(&self.runtime);
        let result = uc
            .components(GraphComponentsInput {
                mode,
                min_size: 1,
                projection: None,
            })
            .map_err(|e| format!("{}", e))?;

        let json = crate::presentation::graph_json::graph_components_json(&result);
        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_graph_cycles(&self, args: &JsonValue) -> Result<String, String> {
        let max_length = args
            .get("max_length")
            .and_then(|v| v.as_u64())
            .map(|v| v as usize)
            .unwrap_or(10);
        let max_cycles = args
            .get("max_cycles")
            .and_then(|v| v.as_u64())
            .map(|v| v as usize)
            .unwrap_or(100);

        let uc = GraphUseCases::new(&self.runtime);
        let result = uc
            .cycles(GraphCyclesInput {
                max_length,
                max_cycles,
                projection: None,
            })
            .map_err(|e| format!("{}", e))?;

        let json = crate::presentation::graph_json::graph_cycles_json(&result);
        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_graph_clustering(&self, _args: &JsonValue) -> Result<String, String> {
        let uc = GraphUseCases::new(&self.runtime);
        let result = uc
            .clustering(GraphClusteringInput {
                top_k: 100,
                include_triangles: true,
                projection: None,
            })
            .map_err(|e| format!("{}", e))?;

        let json = crate::presentation::graph_json::graph_clustering_json(&result);
        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_create_collection(&self, args: &JsonValue) -> Result<String, String> {
        let name = get_str_field(args, "name")?;
        self.runtime
            .db()
            .store()
            .create_collection(name)
            .map_err(|e| format!("{e:?}"))?;
        let mut resp = Map::new();
        resp.insert("ok".into(), JsonValue::Bool(true));
        resp.insert("collection".into(), JsonValue::String(name.to_string()));
        json_to_string(&JsonValue::Object(resp)).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_drop_collection(&self, args: &JsonValue) -> Result<String, String> {
        let name = get_str_field(args, "name")?;
        self.runtime
            .db()
            .store()
            .drop_collection(name)
            .map_err(|e| format!("{e:?}"))?;
        let mut resp = Map::new();
        resp.insert("ok".into(), JsonValue::Bool(true));
        resp.insert("dropped".into(), JsonValue::String(name.to_string()));
        json_to_string(&JsonValue::Object(resp)).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_rql_validate(&self, args: &JsonValue) -> Result<String, String> {
        let rql = get_str_field(args, "rql")?;
        let verdict = reddb_rql::knowledge::validate_rql(rql);
        let json = rql_validation_json(&verdict);
        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_rql_explain(&self, args: &JsonValue) -> Result<String, String> {
        let rql = get_str_field(args, "rql")?;
        let explanation = reddb_rql::knowledge::explain_rql(rql);

        let mut obj = Map::new();
        obj.insert(
            "validation".to_string(),
            rql_validation_json(&explanation.validation),
        );
        if let Some(ast) = &explanation.ast {
            obj.insert("ast".to_string(), JsonValue::String(ast.clone()));
        }
        if let Some(optimized) = &explanation.optimized_ast {
            obj.insert(
                "optimized_ast".to_string(),
                JsonValue::String(optimized.clone()),
            );
        }
        obj.insert(
            "applied_passes".to_string(),
            JsonValue::Array(
                explanation
                    .applied_passes
                    .iter()
                    .cloned()
                    .map(JsonValue::String)
                    .collect(),
            ),
        );
        json_to_string(&JsonValue::Object(obj)).map_err(|e| format!("serialization error: {}", e))
    }

    /// `reddb_type_of`: resolve a type name or a literal value to its canonical
    /// type plus applicable casts/operators, answered entirely from the
    /// generated `reddb-io-types` catalog (ADR 0061). This handler owns no type
    /// knowledge — it only marshals arguments into and JSON out of
    /// [`reddb_types::knowledge::type_of_json`].
    fn tool_type_of(&self, args: &JsonValue) -> Result<String, String> {
        let type_name = args.get("type").and_then(|v| v.as_str());
        let value = args.get("value");
        if type_name.is_none() && value.is_none() {
            return Err(
                "provide one of 'type' (a type name) or 'value' (a JSON literal)".to_string(),
            );
        }
        let facts =
            reddb_types::knowledge::type_of_json(type_name, value).ok_or_else(
                || match type_name {
                    Some(name) => format!("unknown type name '{name}'"),
                    None => "could not resolve a type from the request".to_string(),
                },
            )?;
        json_to_string(&facts).map_err(|e| format!("serialization error: {}", e))
    }

    fn tool_explain_connection(&self, args: &JsonValue) -> Result<String, String> {
        let uri = get_str_field(args, "uri")?;
        let target = reddb_wire::parse(uri).map_err(|err| err.to_string())?;
        let json = connection_explanation_json(uri, &target);
        json_to_string(&json).map_err(|e| format!("serialization error: {}", e))
    }
}

// ------------------------------------------------------------------
// Helpers
// ------------------------------------------------------------------

fn format_mcp_ask_parse_error(err: crate::runtime::ai::mcp_ask_tool::ParseError) -> String {
    use crate::runtime::ai::mcp_ask_tool::ParseError;

    match err {
        ParseError::NotAnObject => "arguments must be an object".to_string(),
        ParseError::MissingQuestion => "missing required field 'question'".to_string(),
        ParseError::QuestionWrongType => "field 'question' must be a string".to_string(),
        ParseError::WrongType { path, expected } => {
            format!("{path} must be {expected}")
        }
        ParseError::OutOfRange { path, detail } => {
            format!("{path} out of range: {detail}")
        }
        ParseError::CacheAndNocache => {
            "options.cache and options.nocache are mutually exclusive".to_string()
        }
        ParseError::UnknownOption { path } => format!("unknown option {path}"),
    }
}

fn connection_explanation_json(uri: &str, target: &reddb_wire::ConnectionTarget) -> JsonValue {
    let scheme = connection_uri_scheme(uri);
    let mut obj = Map::new();
    obj.insert("uri".to_string(), JsonValue::String(uri.to_string()));
    obj.insert("scheme".to_string(), JsonValue::String(scheme.clone()));

    let mut target_obj = Map::new();
    match target {
        reddb_wire::ConnectionTarget::Memory => {
            obj.insert(
                "mode".to_string(),
                JsonValue::String("embedded".to_string()),
            );
            obj.insert(
                "transport".to_string(),
                JsonValue::String("memory".to_string()),
            );
            obj.insert(
                "topology".to_string(),
                JsonValue::String("ephemeral".to_string()),
            );
            target_obj.insert("kind".to_string(), JsonValue::String("memory".to_string()));
        }
        reddb_wire::ConnectionTarget::File { path } => {
            obj.insert(
                "mode".to_string(),
                JsonValue::String("embedded".to_string()),
            );
            obj.insert(
                "transport".to_string(),
                JsonValue::String("filesystem".to_string()),
            );
            obj.insert(
                "topology".to_string(),
                JsonValue::String("persisted".to_string()),
            );
            target_obj.insert("kind".to_string(), JsonValue::String("file".to_string()));
            target_obj.insert(
                "path".to_string(),
                JsonValue::String(path.display().to_string()),
            );
        }
        reddb_wire::ConnectionTarget::Grpc { endpoint } => {
            obj.insert("mode".to_string(), JsonValue::String("remote".to_string()));
            obj.insert(
                "transport".to_string(),
                JsonValue::String(
                    if scheme == "grpcs" {
                        "grpc_tls"
                    } else {
                        "grpc"
                    }
                    .to_string(),
                ),
            );
            obj.insert(
                "topology".to_string(),
                JsonValue::String("single".to_string()),
            );
            target_obj.insert("kind".to_string(), JsonValue::String("grpc".to_string()));
            target_obj.insert("endpoint".to_string(), JsonValue::String(endpoint.clone()));
        }
        reddb_wire::ConnectionTarget::GrpcCluster {
            primary,
            replicas,
            force_primary,
        } => {
            obj.insert("mode".to_string(), JsonValue::String("remote".to_string()));
            obj.insert(
                "transport".to_string(),
                JsonValue::String(cluster_transport_label(&scheme).to_string()),
            );
            obj.insert(
                "topology".to_string(),
                JsonValue::String("primary_replica".to_string()),
            );
            target_obj.insert(
                "kind".to_string(),
                JsonValue::String("primary_replica".to_string()),
            );
            target_obj.insert("primary".to_string(), JsonValue::String(primary.clone()));
            target_obj.insert(
                "replicas".to_string(),
                JsonValue::Array(replicas.iter().cloned().map(JsonValue::String).collect()),
            );
            target_obj.insert("force_primary".to_string(), JsonValue::Bool(*force_primary));
        }
        reddb_wire::ConnectionTarget::Http { base_url } => {
            obj.insert("mode".to_string(), JsonValue::String("remote".to_string()));
            obj.insert(
                "transport".to_string(),
                JsonValue::String(if scheme == "https" { "https" } else { "http" }.to_string()),
            );
            obj.insert(
                "topology".to_string(),
                JsonValue::String("single".to_string()),
            );
            target_obj.insert("kind".to_string(), JsonValue::String("http".to_string()));
            target_obj.insert("base_url".to_string(), JsonValue::String(base_url.clone()));
        }
        reddb_wire::ConnectionTarget::RedWire { host, port, tls } => {
            obj.insert("mode".to_string(), JsonValue::String("remote".to_string()));
            obj.insert(
                "transport".to_string(),
                JsonValue::String(if *tls { "redwire_tls" } else { "redwire" }.to_string()),
            );
            obj.insert(
                "topology".to_string(),
                JsonValue::String("single".to_string()),
            );
            target_obj.insert("kind".to_string(), JsonValue::String("redwire".to_string()));
            target_obj.insert("host".to_string(), JsonValue::String(host.clone()));
            target_obj.insert("port".to_string(), JsonValue::Number(*port as f64));
            target_obj.insert("tls".to_string(), JsonValue::Bool(*tls));
        }
        reddb_wire::ConnectionTarget::WsNative { host, port, tls } => {
            obj.insert("mode".to_string(), JsonValue::String("remote".to_string()));
            obj.insert(
                "transport".to_string(),
                JsonValue::String(if *tls { "redwire_wss" } else { "redwire_ws" }.to_string()),
            );
            obj.insert(
                "topology".to_string(),
                JsonValue::String("single".to_string()),
            );
            target_obj.insert(
                "kind".to_string(),
                JsonValue::String("redwire_websocket".to_string()),
            );
            target_obj.insert("host".to_string(), JsonValue::String(host.clone()));
            target_obj.insert("port".to_string(), JsonValue::Number(*port as f64));
            target_obj.insert("tls".to_string(), JsonValue::Bool(*tls));
        }
    }
    obj.insert("target".to_string(), JsonValue::Object(target_obj));
    JsonValue::Object(obj)
}

fn connection_uri_scheme(uri: &str) -> String {
    uri.split_once(':')
        .map(|(scheme, _)| scheme.to_ascii_lowercase())
        .unwrap_or_default()
}

fn cluster_transport_label(scheme: &str) -> &'static str {
    match scheme {
        "red" => "redwire",
        "reds" => "redwire_tls",
        "grpcs" => "grpc_tls",
        _ => "grpc",
    }
}

/// Render an RQL parse verdict (ADR 0061, #1317) as the JSON the
/// `reddb_rql_validate` / `reddb_rql_explain` tools return.
fn rql_validation_json(verdict: &reddb_rql::knowledge::RqlValidation) -> JsonValue {
    let mut obj = Map::new();
    obj.insert("valid".to_string(), JsonValue::Bool(verdict.valid));
    if let Some(statement) = &verdict.statement {
        obj.insert(
            "statement".to_string(),
            JsonValue::String(statement.clone()),
        );
    }
    obj.insert(
        "has_with_clause".to_string(),
        JsonValue::Bool(verdict.has_with_clause),
    );
    if let Some(error) = &verdict.error {
        obj.insert("error".to_string(), rql_diagnostic_json(error));
    }
    JsonValue::Object(obj)
}

/// Render a structured RQL parse diagnostic as JSON.
fn rql_diagnostic_json(diag: &reddb_rql::knowledge::RqlDiagnostic) -> JsonValue {
    let mut obj = Map::new();
    obj.insert(
        "message".to_string(),
        JsonValue::String(diag.message.clone()),
    );
    obj.insert("line".to_string(), JsonValue::Number(diag.line as f64));
    obj.insert("column".to_string(), JsonValue::Number(diag.column as f64));
    obj.insert("offset".to_string(), JsonValue::Number(diag.offset as f64));
    obj.insert("kind".to_string(), JsonValue::String(diag.kind.clone()));
    obj.insert(
        "expected".to_string(),
        JsonValue::Array(
            diag.expected
                .iter()
                .cloned()
                .map(JsonValue::String)
                .collect(),
        ),
    );
    JsonValue::Object(obj)
}

fn parse_direction(s: Option<&str>) -> RuntimeGraphDirection {
    match s {
        Some("incoming") => RuntimeGraphDirection::Incoming,
        Some("both") => RuntimeGraphDirection::Both,
        _ => RuntimeGraphDirection::Outgoing,
    }
}

/// Parse optional metadata from an `args` JSON object.
fn parse_metadata_arg(
    args: &JsonValue,
) -> Result<Vec<(String, crate::storage::unified::MetadataValue)>, String> {
    match args.get("metadata").and_then(|v| v.as_object()) {
        Some(obj) => {
            let mut out = Vec::with_capacity(obj.len());
            for (key, value) in obj {
                let mv = crate::application::entity::json_to_metadata_value(value)
                    .map_err(|e| format!("{}", e))?;
                out.push((key.clone(), mv));
            }
            Ok(out)
        }
        None => Ok(vec![]),
    }
}

fn parse_string_array_arg(args: &JsonValue, field: &str) -> Result<Vec<String>, String> {
    match args.get(field) {
        None | Some(JsonValue::Null) => Ok(Vec::new()),
        Some(JsonValue::Array(values)) => values
            .iter()
            .map(|value| {
                value
                    .as_str()
                    .map(ToOwned::to_owned)
                    .ok_or_else(|| format!("field '{field}' must be an array of strings"))
            })
            .collect(),
        _ => Err(format!("field '{field}' must be an array of strings")),
    }
}

fn mcp_keyed_ident(value: &str) -> Result<String, String> {
    if !value.is_empty()
        && value
            .bytes()
            .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'.')
    {
        Ok(value.to_string())
    } else {
        Err(
            "keyed collection and key names must use letters, numbers, underscores, or dots"
                .to_string(),
        )
    }
}

fn mcp_value_literal(value: &JsonValue) -> Result<String, String> {
    match value {
        JsonValue::String(value) => Ok(format!("'{}'", value.replace('\'', "''"))),
        JsonValue::Integer(value) => Ok(value.to_string()),
        JsonValue::Number(value) => Ok(value.to_string()),
        JsonValue::Decimal(value) => Ok(value.clone()),
        JsonValue::Bool(value) => Ok(value.to_string()),
        JsonValue::Null => Ok("NULL".to_string()),
        JsonValue::Array(_) | JsonValue::Object(_) => {
            json_to_string(value).map_err(|err| format!("serialization error: {err}"))
        }
    }
}

fn append_mcp_tags_clause(sql: &mut String, tags: &[String]) {
    if tags.is_empty() {
        return;
    }
    sql.push_str(" TAGS [");
    for (index, tag) in tags.iter().enumerate() {
        if index > 0 {
            sql.push_str(", ");
        }
        sql.push('\'');
        sql.push_str(&tag.replace('\'', "''"));
        sql.push('\'');
    }
    sql.push(']');
}

fn reject_mcp_volatile_options(args: &JsonValue, domain: &str) -> Result<(), String> {
    for field in ["ttl", "ttl_ms", "expire", "expire_ms", "expires_at"] {
        if args.get(field).is_some() {
            return Err(format!(
                "{domain} does not support TTL or expiration options"
            ));
        }
    }
    Ok(())
}

// Convert a storage Value to JSON (local helper to avoid visibility issues).
fn get_str_field<'a>(args: &'a JsonValue, field: &str) -> Result<&'a str, String> {
    args.get(field)
        .and_then(|v| v.as_str())
        .ok_or_else(|| format!("missing '{field}'"))
}

// Auth tool implementations
impl McpServer {
    fn tool_auth_bootstrap(&self, args: &JsonValue) -> Result<String, String> {
        let username = get_str_field(args, "username")?;
        let password = get_str_field(args, "password")?;

        let br = self
            .auth_store
            .bootstrap(username, password)
            .map_err(|e| e.to_string())?;

        let mut result = Map::new();
        result.insert("ok".into(), JsonValue::Bool(true));
        result.insert("username".into(), JsonValue::String(br.user.username));
        result.insert(
            "role".into(),
            JsonValue::String(br.user.role.as_str().into()),
        );
        result.insert("api_key".into(), JsonValue::String(br.api_key.key));
        result.insert("api_key_name".into(), JsonValue::String(br.api_key.name));
        if let Some(cert) = br.certificate {
            result.insert("certificate".into(), JsonValue::String(cert));
            result.insert(
                "message".into(),
                JsonValue::String(
                    "Save this certificate — it is the ONLY way to unseal the vault after restart."
                        .into(),
                ),
            );
        } else {
            result.insert(
                "message".into(),
                JsonValue::String(
                    "First admin user created. Save the API key — it won't be shown again.".into(),
                ),
            );
        }
        json_to_string(&JsonValue::Object(result))
    }

    fn tool_auth_create_user(&self, args: &JsonValue) -> Result<String, String> {
        let username = get_str_field(args, "username")?;
        let password = get_str_field(args, "password")?;
        let role_str = get_str_field(args, "role")?;
        let role = Role::from_str(role_str).ok_or_else(|| format!("invalid role: {role_str}"))?;

        self.auth_store
            .create_user(username, password, role)
            .map_err(|e| e.to_string())?;

        let mut result = Map::new();
        result.insert("ok".into(), JsonValue::Bool(true));
        result.insert("username".into(), JsonValue::String(username.into()));
        result.insert("role".into(), JsonValue::String(role.as_str().into()));
        json_to_string(&JsonValue::Object(result))
    }

    fn tool_auth_login(&self, args: &JsonValue) -> Result<String, String> {
        let username = get_str_field(args, "username")?;
        let password = get_str_field(args, "password")?;

        let session = self
            .auth_store
            .authenticate_active(None, username, password)
            .map_err(|e| e.to_string())?;

        let mut result = Map::new();
        result.insert("ok".into(), JsonValue::Bool(true));
        result.insert("token".into(), JsonValue::String(session.token));
        result.insert("username".into(), JsonValue::String(session.username));
        result.insert(
            "role".into(),
            JsonValue::String(session.role.as_str().into()),
        );
        result.insert(
            "expires_at".into(),
            JsonValue::Number(session.expires_at as f64),
        );
        json_to_string(&JsonValue::Object(result))
    }

    fn tool_auth_create_api_key(&self, args: &JsonValue) -> Result<String, String> {
        let username = get_str_field(args, "username")?;
        let name = get_str_field(args, "name")?;
        let role_str = get_str_field(args, "role")?;
        let role = Role::from_str(role_str).ok_or_else(|| format!("invalid role: {role_str}"))?;

        let key = self
            .auth_store
            .create_api_key(username, name, role)
            .map_err(|e| e.to_string())?;

        let mut result = Map::new();
        result.insert("ok".into(), JsonValue::Bool(true));
        result.insert("key".into(), JsonValue::String(key.key));
        result.insert("name".into(), JsonValue::String(key.name));
        result.insert("role".into(), JsonValue::String(key.role.as_str().into()));
        json_to_string(&JsonValue::Object(result))
    }

    fn tool_auth_list_users(&self) -> Result<String, String> {
        let users = self.auth_store.list_users();
        let arr: Vec<JsonValue> = users
            .into_iter()
            .map(|u| {
                let mut obj = Map::new();
                obj.insert("username".into(), JsonValue::String(u.username));
                obj.insert("role".into(), JsonValue::String(u.role.as_str().into()));
                obj.insert("enabled".into(), JsonValue::Bool(u.enabled));
                obj.insert(
                    "api_key_count".into(),
                    JsonValue::Number(u.api_keys.len() as f64),
                );
                JsonValue::Object(obj)
            })
            .collect();
        json_to_string(&JsonValue::Array(arr))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::{Read, Write};
    use std::net::{SocketAddr, TcpListener, TcpStream};
    use std::sync::atomic::{AtomicBool, Ordering};
    use std::sync::{Arc, Mutex};
    use std::thread::{self, JoinHandle};
    use std::time::Duration;

    static ASK_ENV_LOCK: Mutex<()> = Mutex::new(());

    /// The existing suite exercises the whole tool surface (DDL, inserts,
    /// vault, auth), so it runs at the level an operator would opt into with
    /// `--allow-admin`. `read_only_session_*` below cover the default.
    fn make_server() -> McpServer {
        let rt = RedDBRuntime::in_memory().expect("in-memory runtime");
        McpServer::new(rt).with_capability(McpCapability::Admin)
    }

    /// A session at the shipped default: no flags.
    fn make_read_only_server() -> McpServer {
        let rt = RedDBRuntime::in_memory().expect("in-memory runtime");
        McpServer::new(rt)
    }

    #[test]
    fn update_tool_renders_values_as_literals_and_refuses_crafted_identifiers() {
        let srv = make_server();
        srv.tool_query(&parse_json(
            r#"{"sql":"CREATE TABLE mcp_upd (id INTEGER, name TEXT)"}"#,
        ))
        .expect("create");
        srv.tool_query(&parse_json(
            r#"{"sql":"INSERT INTO mcp_upd (id, name) VALUES (1, 'a')"}"#,
        ))
        .expect("seed");
        // A quote inside the value used to close the literal and run the
        // remainder as SQL.
        srv.tool_update(&parse_json(
            r#"{"collection":"mcp_upd","set":{"name":"o'brien"},"where_filter":"id = 1"}"#,
        ))
        .expect("escaped value");
        let rows = srv
            .tool_query(&parse_json(
                r#"{"sql":"SELECT name FROM mcp_upd WHERE id = 1"}"#,
            ))
            .expect("read back");
        assert!(rows.contains("o'brien"), "{rows}");

        for args in [
            r#"{"collection":"mcp_upd; DROP TABLE mcp_upd","set":{"name":"x"}}"#,
            r#"{"collection":"mcp_upd","set":{"name = 'x' WHERE 1=1 --":"x"}}"#,
        ] {
            let err = srv
                .tool_update(&parse_json(args))
                .expect_err("crafted identifier must be refused");
            assert!(err.contains("invalid"), "{err}");
        }
    }

    #[test]
    fn read_only_session_refuses_mutating_tools() {
        // The MCP transport has no second principal to authenticate, so the
        // tool list is the boundary: a model steered by injected content
        // must not reach `drop_collection` or the vault by default.
        let srv = make_read_only_server();
        assert_eq!(srv.capability(), McpCapability::ReadOnly);
        for tool in [
            "reddb_drop_collection",
            "reddb_delete",
            "reddb_insert_row",
            "reddb_vault_unseal",
            "reddb_auth_create_api_key",
        ] {
            let err = srv
                .check_tool_capability(tool)
                .expect_err("mutating tool must be refused by default");
            assert!(err.contains("read-only"), "{tool}: {err}");
        }
        // Reads still work, otherwise the default is useless.
        srv.check_tool_capability("reddb_query")
            .expect("reads stay available");
        srv.check_tool_capability("reddb_collections")
            .expect("introspection stays available");
    }

    #[test]
    fn read_only_session_refuses_mutating_sql_through_the_query_tool() {
        // `reddb_query` is a read tool that takes arbitrary SQL, so the
        // statement has to be classified too — otherwise the one tool a
        // read-only session may call is a full write channel.
        let srv = make_read_only_server();
        let err = srv
            .tool_query(&parse_json(r#"{"sql":"DROP TABLE users"}"#))
            .expect_err("DDL must be refused in a read-only session");
        assert!(err.contains("write"), "{err}");

        let err = srv
            .tool_query(&parse_json(r#"{"sql":"SHOW SECRETS"}"#))
            .expect_err("secret reads must require admin");
        assert!(err.contains("admin"), "{err}");

        srv.tool_query(&parse_json(r#"{"sql":"SELECT 1 AS one"}"#))
            .expect("plain reads still work");
    }

    #[test]
    fn allow_write_does_not_unlock_vault_or_auth_tools() {
        // `--allow-write` is for data and schema. Credentials need the
        // separate, louder opt-in.
        let rt = RedDBRuntime::in_memory().expect("in-memory runtime");
        let srv = McpServer::new(rt).with_capability(McpCapability::Write);
        srv.check_tool_capability("reddb_insert_row")
            .expect("writes are enabled");
        let err = srv
            .check_tool_capability("reddb_vault_get")
            .expect_err("vault stays behind --allow-admin");
        assert!(err.contains("admin"), "{err}");
    }

    fn parse_json(s: &str) -> JsonValue {
        json_from_str(s).expect("valid json")
    }

    #[test]
    fn initialize_advertises_resources_capability() {
        let mut srv = make_server();
        let response = srv.handle_initialize(Some(&JsonValue::Number(1.0)));
        let parsed = parse_json(&response);
        let resources = parsed
            .get("result")
            .and_then(|result| result.get("capabilities"))
            .and_then(|caps| caps.get("resources"))
            .expect("resources capability advertised");
        assert_eq!(
            resources.get("subscribe").and_then(JsonValue::as_bool),
            Some(false)
        );
    }

    #[test]
    fn resources_list_includes_rql_knowledge() {
        let mut srv = make_server();
        let request = parse_json(r#"{"jsonrpc":"2.0","id":7,"method":"resources/list"}"#);
        let response = srv.handle_message(&request).expect("response");
        let parsed = parse_json(&response);
        let resources = parsed
            .get("result")
            .and_then(|result| result.get("resources"))
            .and_then(JsonValue::as_array)
            .expect("resources array");

        let rql = resources
            .iter()
            .find(|res| res.get("uri").and_then(JsonValue::as_str) == Some("reddb://knowledge/rql"))
            .expect("rql knowledge resource listed");
        assert_eq!(
            rql.get("mimeType").and_then(JsonValue::as_str),
            Some("text/markdown")
        );
        assert!(rql.get("name").and_then(JsonValue::as_str).is_some());
    }

    #[test]
    fn resources_list_includes_connection_knowledge() {
        let mut srv = make_server();
        let request = parse_json(r#"{"jsonrpc":"2.0","id":11,"method":"resources/list"}"#);
        let response = srv.handle_message(&request).expect("response");
        let parsed = parse_json(&response);
        let resources = parsed
            .get("result")
            .and_then(|result| result.get("resources"))
            .and_then(JsonValue::as_array)
            .expect("resources array");

        let connection = resources
            .iter()
            .find(|res| {
                res.get("uri").and_then(JsonValue::as_str)
                    == Some(reddb_wire::knowledge::RESOURCE_URI)
            })
            .expect("connection knowledge resource listed");
        assert_eq!(
            connection.get("mimeType").and_then(JsonValue::as_str),
            Some("text/markdown")
        );
        assert!(connection.get("name").and_then(JsonValue::as_str).is_some());
    }

    #[test]
    fn resources_read_returns_generated_rql_reference() {
        let mut srv = make_server();
        let request = parse_json(
            r#"{"jsonrpc":"2.0","id":8,"method":"resources/read","params":{"uri":"reddb://knowledge/rql"}}"#,
        );
        let response = srv.handle_message(&request).expect("response");
        let parsed = parse_json(&response);
        let text = parsed
            .get("result")
            .and_then(|result| result.get("contents"))
            .and_then(JsonValue::as_array)
            .and_then(|contents| contents.first())
            .and_then(|item| item.get("text"))
            .and_then(JsonValue::as_str)
            .expect("resource text");

        // The body is exactly what the RQL authority generates — same source as
        // docs/llms.txt.
        assert_eq!(text, reddb_rql::knowledge::rql_reference_markdown());
        // And it is genuinely generated from the engine: a sampled keyword and
        // function from the authorities are present.
        assert!(text.contains("`SELECT`"), "keyword missing: {text:.120}");
        assert!(text.contains("`COUNT`"), "function missing");
    }

    #[test]
    fn resources_read_returns_generated_type_reference() {
        let mut srv = make_server();
        let request = parse_json(
            r#"{"jsonrpc":"2.0","id":10,"method":"resources/read","params":{"uri":"reddb://knowledge/types"}}"#,
        );
        let response = srv.handle_message(&request).expect("response");
        let parsed = parse_json(&response);
        let text = parsed
            .get("result")
            .and_then(|result| result.get("contents"))
            .and_then(JsonValue::as_array)
            .and_then(|contents| contents.first())
            .and_then(|item| item.get("text"))
            .and_then(JsonValue::as_str)
            .expect("resource text");

        // The body is exactly what the type authority generates — same source as
        // docs/llms.txt.
        assert_eq!(text, reddb_types::knowledge::type_reference_markdown());
        // And it is genuinely generated from the engine: a sampled value type and
        // a multi-model paradigm are present.
        assert!(
            text.contains("`INTEGER`"),
            "value type missing: {text:.120}"
        );
        assert!(text.contains("Multi-model map"), "multi-model map missing");
    }

    #[test]
    fn resources_read_returns_generated_connection_reference() {
        let mut srv = make_server();
        let request = parse_json(
            r#"{"jsonrpc":"2.0","id":12,"method":"resources/read","params":{"uri":"reddb://knowledge/connections"}}"#,
        );
        let response = srv.handle_message(&request).expect("response");
        let parsed = parse_json(&response);
        let text = parsed
            .get("result")
            .and_then(|result| result.get("contents"))
            .and_then(JsonValue::as_array)
            .and_then(|contents| contents.first())
            .and_then(|item| item.get("text"))
            .and_then(JsonValue::as_str)
            .expect("resource text");

        // The body is exactly what the wire authority generates — same source as
        // docs/llms.txt.
        assert_eq!(text, reddb_wire::knowledge::connection_reference_markdown());
        assert!(text.contains("`red://`"), "red scheme missing: {text:.120}");
        assert!(text.contains("`reds://`"), "reds scheme missing");
        assert!(
            text.contains("principal transport"),
            "principal transport narrative missing"
        );
        assert!(
            text.contains("Embedded / standalone"),
            "embedded topology missing"
        );
        assert!(text.contains("Serverless"), "serverless topology missing");
        assert!(
            text.contains("Primary-replica"),
            "primary-replica topology missing"
        );
        assert!(text.contains("Cluster"), "cluster topology missing");
    }

    #[test]
    fn resources_read_unknown_uri_errors() {
        let mut srv = make_server();
        let request = parse_json(
            r#"{"jsonrpc":"2.0","id":9,"method":"resources/read","params":{"uri":"reddb://knowledge/nope"}}"#,
        );
        let response = srv.handle_message(&request).expect("response");
        let parsed = parse_json(&response);
        assert!(
            parsed.get("error").is_some(),
            "unknown resource must error: {response}"
        );
    }

    #[test]
    fn tools_list_registers_reddb_ask_descriptor() {
        let srv = make_server();
        let response = srv.handle_tools_list(Some(&JsonValue::Number(1.0)));
        let parsed = parse_json(&response);
        let tools = parsed
            .get("result")
            .and_then(|result| result.get("tools"))
            .and_then(JsonValue::as_array)
            .expect("tools array");

        let ask = tools
            .iter()
            .find(|tool| tool.get("name").and_then(JsonValue::as_str) == Some("reddb.ask"))
            .expect("reddb.ask registered");

        let desc = ask
            .get("description")
            .and_then(JsonValue::as_str)
            .expect("description");
        assert!(desc.contains("citations"), "description: {desc}");
        assert!(desc.contains("sources_flat"), "description: {desc}");
        assert!(desc.contains("URN"), "description: {desc}");

        let options = ask
            .get("inputSchema")
            .and_then(|schema| schema.get("properties"))
            .and_then(|props| props.get("options"))
            .and_then(|opts| opts.get("properties"))
            .and_then(JsonValue::as_object)
            .expect("options properties");
        for key in [
            "strict",
            "using",
            "model",
            "limit",
            "min_score",
            "depth",
            "temperature",
            "seed",
            "cache",
            "nocache",
        ] {
            assert!(
                options.contains_key(key),
                "missing option {key} in {options:?}"
            );
        }
    }

    #[test]
    fn embedded_dispatch_reaches_every_advertised_tool() {
        let server = make_server();
        let mut tool_names: Vec<&str> = tools::all_tools()
            .into_iter()
            .map(|tool| tool.name)
            .collect();
        tool_names.push(crate::runtime::ai::mcp_ask_tool::TOOL_NAME);

        for tool_name in tool_names {
            let mut params = Map::new();
            params.insert("name".to_string(), JsonValue::String(tool_name.to_string()));
            params.insert("arguments".to_string(), JsonValue::Object(Map::new()));
            let response = server.handle_tools_call(
                Some(&JsonValue::Number(1.0)),
                Some(&JsonValue::Object(params)),
            );
            let parsed = parse_json(&response);
            let result = parsed.get("result").expect("tool result");
            let text = result
                .get("content")
                .and_then(JsonValue::as_array)
                .and_then(|content| content.first())
                .and_then(|item| item.get("text"))
                .and_then(JsonValue::as_str)
                .expect("tool result text");

            assert!(
                !text.contains("unknown tool") && !text.contains("denied by policy"),
                "embedded tool {tool_name} did not reach its dispatch arm: {text}"
            );
        }
    }

    #[test]
    fn tools_call_reddb_ask_uses_typed_argument_parser() {
        let srv = make_server();
        let params = parse_json(
            r#"{
                "name": "reddb.ask",
                "arguments": {
                    "question": "what cites this?",
                    "options": { "tempurature": 0.2 }
                }
            }"#,
        );

        let response = srv.handle_tools_call(Some(&JsonValue::Number(1.0)), Some(&params));
        let parsed = parse_json(&response);
        let result = parsed.get("result").expect("result");
        assert_eq!(
            result.get("isError").and_then(JsonValue::as_bool),
            Some(true)
        );
        let text = result
            .get("content")
            .and_then(JsonValue::as_array)
            .and_then(|content| content.first())
            .and_then(|item| item.get("text"))
            .and_then(JsonValue::as_str)
            .expect("error text");
        assert!(text.contains("options.tempurature"), "text: {text}");
    }

    /// Helper: drive a `reddb_type_of` `tools/call` over the JSON-RPC seam and
    /// return the parsed result object (`content`/`isError`), exercising the
    /// real dispatcher.
    fn call_type_of(srv: &mut McpServer, arguments: &str) -> JsonValue {
        let request = parse_json(&format!(
            r#"{{"jsonrpc":"2.0","id":21,"method":"tools/call","params":{{"name":"reddb_type_of","arguments":{arguments}}}}}"#
        ));
        let response = srv.handle_message(&request).expect("response");
        parse_json(&response)
            .get("result")
            .cloned()
            .expect("result object")
    }

    /// Extract the text payload of a (non-error) tool result and parse it as JSON.
    fn type_of_result_json(result: &JsonValue) -> JsonValue {
        assert_ne!(
            result.get("isError").and_then(JsonValue::as_bool),
            Some(true),
            "tool returned an error: {result:?}"
        );
        let text = result
            .get("content")
            .and_then(JsonValue::as_array)
            .and_then(|content| content.first())
            .and_then(|item| item.get("text"))
            .and_then(JsonValue::as_str)
            .expect("result text");
        parse_json(text)
    }

    #[test]
    fn tools_list_includes_reddb_type_of() {
        let srv = make_server();
        let response = srv.handle_tools_list(Some(&JsonValue::Number(1.0)));
        let parsed = parse_json(&response);
        let tools = parsed
            .get("result")
            .and_then(|result| result.get("tools"))
            .and_then(JsonValue::as_array)
            .expect("tools array");
        assert!(
            tools
                .iter()
                .any(|tool| tool.get("name").and_then(JsonValue::as_str) == Some("reddb_type_of")),
            "reddb_type_of must be advertised in tools/list"
        );
    }

    #[test]
    fn tools_call_type_of_resolves_type_name_from_catalog() {
        let mut srv = make_server();
        let result = call_type_of(&mut srv, r#"{"type":"int"}"#);
        let facts = type_of_result_json(&result);

        // Canonical type + category come from the catalog authority.
        assert_eq!(
            facts.get("canonical_type").and_then(JsonValue::as_str),
            Some("INTEGER")
        );
        assert_eq!(
            facts.get("category").and_then(JsonValue::as_str),
            Some("Numeric")
        );
        // Applicable casts include the implicit INTEGER → FLOAT widening.
        let casts = facts.get("casts").and_then(JsonValue::as_array).unwrap();
        assert!(
            casts.iter().any(|c| {
                c.get("target").and_then(JsonValue::as_str) == Some("FLOAT")
                    && c.get("context").and_then(JsonValue::as_str) == Some("implicit")
            }),
            "casts: {casts:?}"
        );
        // Applicable operators include arithmetic +.
        let operators = facts
            .get("operators")
            .and_then(JsonValue::as_array)
            .unwrap();
        assert!(
            operators
                .iter()
                .any(|o| o.get("symbol").and_then(JsonValue::as_str) == Some("+")),
            "operators: {operators:?}"
        );
    }

    #[test]
    fn tools_call_type_of_infers_type_from_literal_value() {
        let mut srv = make_server();
        let result = call_type_of(&mut srv, r#"{"value":true}"#);
        let facts = type_of_result_json(&result);
        assert_eq!(
            facts.get("canonical_type").and_then(JsonValue::as_str),
            Some("BOOLEAN")
        );
    }

    #[test]
    fn tools_call_type_of_unknown_name_is_error() {
        let mut srv = make_server();
        let result = call_type_of(&mut srv, r#"{"type":"frobnicate"}"#);
        assert_eq!(
            result.get("isError").and_then(JsonValue::as_bool),
            Some(true)
        );
        let text = result
            .get("content")
            .and_then(JsonValue::as_array)
            .and_then(|content| content.first())
            .and_then(|item| item.get("text"))
            .and_then(JsonValue::as_str)
            .expect("error text");
        assert!(text.contains("frobnicate"), "text: {text}");
    }

    #[test]
    fn tools_call_type_of_requires_an_argument() {
        let mut srv = make_server();
        let result = call_type_of(&mut srv, r#"{}"#);
        assert_eq!(
            result.get("isError").and_then(JsonValue::as_bool),
            Some(true)
        );
    }

    #[test]
    fn tools_call_reddb_ask_returns_canonical_citation_envelope() {
        let _guard = ASK_ENV_LOCK.lock().expect("env lock");
        let stub = AskStub::start();
        let _api_base = EnvVarGuard::set(
            "REDDB_OLLAMA_API_BASE",
            &format!("http://{}/v1", stub.addr()),
        );

        let srv = make_server();
        srv.tool_query(&parse_json(
            r#"{"sql":"CREATE TABLE travel (id INTEGER, passport TEXT, notes TEXT)"}"#,
        ))
        .expect("ddl ok");
        srv.tool_query(&parse_json(
            r#"{"sql":"INSERT INTO travel (id, passport, notes) VALUES (1, 'PT-002', 'incident FDD-12313 escalated')"}"#,
        ))
        .expect("insert ok");

        let params = parse_json(
            r#"{
                "name": "reddb.ask",
                "arguments": {
                    "question": "passport FDD-12313",
                    "options": {
                        "strict": false,
                        "using": "ollama",
                        "model": "mock-ask",
                        "limit": 1,
                        "min_score": 0,
                        "depth": 0,
                        "temperature": 0,
                        "seed": 0,
                        "cache": { "ttl": "5m" }
                    }
                }
            }"#,
        );

        let response = srv.handle_tools_call(Some(&JsonValue::Number(1.0)), Some(&params));
        let parsed = parse_json(&response);
        let result = parsed.get("result").expect("result");
        assert_ne!(
            result.get("isError").and_then(JsonValue::as_bool),
            Some(true),
            "response: {response}"
        );
        let text = result
            .get("content")
            .and_then(JsonValue::as_array)
            .and_then(|content| content.first())
            .and_then(|item| item.get("text"))
            .and_then(JsonValue::as_str)
            .expect("tool text");
        let envelope = parse_json(text);

        assert_eq!(
            envelope.get("answer").and_then(JsonValue::as_str),
            Some("FDD-12313 escalated [^1].")
        );
        assert_eq!(
            envelope.get("provider").and_then(JsonValue::as_str),
            Some("ollama")
        );
        assert_eq!(
            envelope.get("model").and_then(JsonValue::as_str),
            Some("mock-ask")
        );
        assert_eq!(
            envelope.get("cache_hit").and_then(JsonValue::as_bool),
            Some(false)
        );
        assert!(envelope
            .get("sources_flat")
            .and_then(JsonValue::as_array)
            .is_some());
        assert!(envelope
            .get("citations")
            .and_then(JsonValue::as_array)
            .is_some());
        assert!(envelope
            .get("validation")
            .and_then(JsonValue::as_object)
            .is_some());
        assert!(
            envelope.get("rows").is_none(),
            "ASK must not be row-wrapped: {text}"
        );
    }

    #[test]
    fn tool_query_without_params_keeps_legacy_path() {
        let srv = make_server();
        let args = parse_json(r#"{"sql":"SELECT 1 AS one"}"#);
        let out = srv.tool_query(&args).expect("query ok");
        assert!(out.contains("\"one\""), "expected 'one' column in {out}");
    }

    #[test]
    fn tool_query_binds_int_and_text_params() {
        let srv = make_server();
        srv.tool_query(&parse_json(
            r#"{"sql":"CREATE TABLE mcpp (id INTEGER, name TEXT)"}"#,
        ))
        .expect("ddl ok");
        srv.tool_query(&parse_json(
            r#"{"sql":"INSERT INTO mcpp (id, name) VALUES (1, 'Alice')"}"#,
        ))
        .expect("insert 1");
        srv.tool_query(&parse_json(
            r#"{"sql":"INSERT INTO mcpp (id, name) VALUES (2, 'Bob')"}"#,
        ))
        .expect("insert 2");

        let out = srv
            .tool_query(&parse_json(
                r#"{"sql":"SELECT * FROM mcpp WHERE id = $1 AND name = $2","params":[1,"Alice"]}"#,
            ))
            .expect("query with params ok");
        assert!(out.contains("Alice"), "expected Alice in {out}");
        assert!(!out.contains("Bob"), "Bob must not match: {out}");
    }

    #[test]
    fn tool_query_params_must_be_array() {
        let srv = make_server();
        let err = srv
            .tool_query(&parse_json(
                r#"{"sql":"SELECT 1","params":{"not":"array"}}"#,
            ))
            .expect_err("must reject non-array params");
        assert!(err.contains("array"), "got {err}");
    }

    #[test]
    fn tool_query_param_arity_mismatch_surfaces_error() {
        let srv = make_server();
        srv.tool_query(&parse_json(r#"{"sql":"CREATE TABLE mcpa (id INTEGER)"}"#))
            .expect("ddl ok");
        let err = srv
            .tool_query(&parse_json(
                r#"{"sql":"SELECT * FROM mcpa WHERE id = $1","params":[1,2]}"#,
            ))
            .expect_err("arity mismatch");
        assert!(
            err.contains("number of parameters") || err.contains("expects"),
            "got {err}"
        );
    }

    #[test]
    fn tool_query_vector_param_binds_into_search_similar() {
        let srv = make_server();
        let out = srv.tool_query(&parse_json(
            r#"{"sql":"SEARCH SIMILAR $1 COLLECTION mcpv LIMIT 5","params":[[0.1,0.2,0.3]]}"#,
        ));
        // The collection doesn't exist; we only need to confirm the
        // param-bind path runs (i.e. the error reflects runtime semantics,
        // not a `$N` placeholder being unresolved).
        if let Err(e) = out {
            assert!(
                !e.contains("placeholder") && !e.contains("Parameter"),
                "param did not bind: {e}"
            );
        }
    }

    // ------------------------------------------------------------------
    // RQL validate / explain tools (ADR 0061, #1317), exercised at the
    // JSON-RPC tools/call seam.
    // ------------------------------------------------------------------

    fn call_tool(srv: &McpServer, name: &str, rql: &str) -> JsonValue {
        let params = parse_json(
            &json_to_string(&{
                let mut p = Map::new();
                p.insert("name".to_string(), JsonValue::String(name.to_string()));
                let mut args = Map::new();
                args.insert("rql".to_string(), JsonValue::String(rql.to_string()));
                p.insert("arguments".to_string(), JsonValue::Object(args));
                JsonValue::Object(p)
            })
            .expect("serialize params"),
        );
        let response = srv.handle_tools_call(Some(&JsonValue::Number(1.0)), Some(&params));
        parse_json(&response)
    }

    fn tool_result_text(parsed: &JsonValue) -> String {
        parsed
            .get("result")
            .and_then(|r| r.get("content"))
            .and_then(JsonValue::as_array)
            .and_then(|content| content.first())
            .and_then(|item| item.get("text"))
            .and_then(JsonValue::as_str)
            .expect("tool text")
            .to_string()
    }

    #[test]
    fn tools_list_registers_rql_validate_and_explain() {
        let srv = make_server();
        let response = srv.handle_tools_list(Some(&JsonValue::Number(1.0)));
        let parsed = parse_json(&response);
        let tools = parsed
            .get("result")
            .and_then(|result| result.get("tools"))
            .and_then(JsonValue::as_array)
            .expect("tools array");
        let names: Vec<&str> = tools
            .iter()
            .filter_map(|t| t.get("name").and_then(JsonValue::as_str))
            .collect();
        assert!(names.contains(&"reddb_rql_validate"), "{names:?}");
        assert!(names.contains(&"reddb_rql_explain"), "{names:?}");
    }

    #[test]
    fn tools_call_rql_validate_accepts_valid_query() {
        let srv = make_server();
        let parsed = call_tool(&srv, "reddb_rql_validate", "SELECT * FROM users");
        let result = parsed.get("result").expect("result");
        assert_ne!(
            result.get("isError").and_then(JsonValue::as_bool),
            Some(true)
        );
        let envelope = parse_json(&tool_result_text(&parsed));
        assert_eq!(
            envelope.get("valid").and_then(JsonValue::as_bool),
            Some(true)
        );
        assert_eq!(
            envelope.get("statement").and_then(JsonValue::as_str),
            Some("Table")
        );
    }

    #[test]
    fn tools_call_rql_validate_reports_structured_error() {
        let srv = make_server();
        let parsed = call_tool(&srv, "reddb_rql_validate", "SELECT * FROM");
        let envelope = parse_json(&tool_result_text(&parsed));
        assert_eq!(
            envelope.get("valid").and_then(JsonValue::as_bool),
            Some(false)
        );
        let error = envelope.get("error").expect("structured error");
        assert!(error
            .get("message")
            .and_then(JsonValue::as_str)
            .is_some_and(|m| !m.is_empty()));
        assert!(error.get("line").and_then(JsonValue::as_f64).is_some());
        assert!(error
            .get("kind")
            .and_then(JsonValue::as_str)
            .is_some_and(|k| !k.is_empty()));
    }

    #[test]
    fn tools_call_rql_explain_returns_ast_and_plan() {
        let srv = make_server();
        let parsed = call_tool(
            &srv,
            "reddb_rql_explain",
            "SELECT * FROM users WHERE id = 1",
        );
        let envelope = parse_json(&tool_result_text(&parsed));
        assert_eq!(
            envelope
                .get("validation")
                .and_then(|v| v.get("valid"))
                .and_then(JsonValue::as_bool),
            Some(true)
        );
        assert!(envelope.get("ast").and_then(JsonValue::as_str).is_some());
        assert!(envelope
            .get("optimized_ast")
            .and_then(JsonValue::as_str)
            .is_some());
        assert!(envelope
            .get("applied_passes")
            .and_then(JsonValue::as_array)
            .is_some());
    }

    #[test]
    fn tools_call_rql_validate_missing_field_errors() {
        let srv = make_server();
        let params = parse_json(r#"{"name":"reddb_rql_validate","arguments":{}}"#);
        let response = srv.handle_tools_call(Some(&JsonValue::Number(1.0)), Some(&params));
        let parsed = parse_json(&response);
        let result = parsed.get("result").expect("result");
        assert_eq!(
            result.get("isError").and_then(JsonValue::as_bool),
            Some(true)
        );
    }

    fn call_explain_connection(srv: &mut McpServer, uri: &str) -> JsonValue {
        let request = parse_json(
            &json_to_string(&{
                let mut request = Map::new();
                request.insert("jsonrpc".to_string(), JsonValue::String("2.0".to_string()));
                request.insert("id".to_string(), JsonValue::Number(31.0));
                request.insert(
                    "method".to_string(),
                    JsonValue::String("tools/call".to_string()),
                );

                let mut args = Map::new();
                args.insert("uri".to_string(), JsonValue::String(uri.to_string()));

                let mut params = Map::new();
                params.insert(
                    "name".to_string(),
                    JsonValue::String("reddb_explain_connection".to_string()),
                );
                params.insert("arguments".to_string(), JsonValue::Object(args));
                request.insert("params".to_string(), JsonValue::Object(params));
                JsonValue::Object(request)
            })
            .expect("serialize request"),
        );
        let response = srv.handle_message(&request).expect("response");
        let parsed = parse_json(&response);
        let result = parsed.get("result").expect("result");
        assert_ne!(
            result.get("isError").and_then(JsonValue::as_bool),
            Some(true),
            "tool returned error: {parsed:?}"
        );
        parse_json(&tool_result_text(&parsed))
    }

    #[test]
    fn tools_call_explain_connection_maps_supported_schemes() {
        let mut srv = make_server();
        for (uri, mode, transport, topology) in [
            ("memory://", "embedded", "memory", "ephemeral"),
            (
                "file:///tmp/reddb/data.rdb",
                "embedded",
                "filesystem",
                "persisted",
            ),
            ("red://db.example:5050", "remote", "redwire", "single"),
            ("reds://db.example", "remote", "redwire_tls", "single"),
            ("grpc://db.example", "remote", "grpc", "single"),
            ("http://db.example", "remote", "http", "single"),
            ("https://db.example", "remote", "https", "single"),
        ] {
            let envelope = call_explain_connection(&mut srv, uri);
            assert_eq!(envelope.get("uri").and_then(JsonValue::as_str), Some(uri));
            assert_eq!(
                envelope.get("mode").and_then(JsonValue::as_str),
                Some(mode),
                "{uri}: {envelope:?}"
            );
            assert_eq!(
                envelope.get("transport").and_then(JsonValue::as_str),
                Some(transport),
                "{uri}: {envelope:?}"
            );
            assert_eq!(
                envelope.get("topology").and_then(JsonValue::as_str),
                Some(topology),
                "{uri}: {envelope:?}"
            );
        }

        let cluster = call_explain_connection(&mut srv, "grpc://primary,replica?route=primary");
        assert_eq!(
            cluster.get("topology").and_then(JsonValue::as_str),
            Some("primary_replica")
        );
        assert_eq!(
            cluster.get("transport").and_then(JsonValue::as_str),
            Some("grpc")
        );
        assert_eq!(
            cluster
                .get("target")
                .and_then(|target| target.get("force_primary"))
                .and_then(JsonValue::as_bool),
            Some(true)
        );
    }

    struct EnvVarGuard {
        name: &'static str,
        previous: Option<String>,
    }

    impl EnvVarGuard {
        fn set(name: &'static str, value: &str) -> Self {
            let previous = std::env::var(name).ok();
            std::env::set_var(name, value);
            Self { name, previous }
        }
    }

    impl Drop for EnvVarGuard {
        fn drop(&mut self) {
            if let Some(value) = self.previous.take() {
                std::env::set_var(self.name, value);
            } else {
                std::env::remove_var(self.name);
            }
        }
    }

    struct AskStub {
        addr: SocketAddr,
        shutdown: Arc<AtomicBool>,
        handle: Option<JoinHandle<()>>,
    }

    impl AskStub {
        fn start() -> Self {
            let listener = TcpListener::bind("127.0.0.1:0").expect("stub bind");
            listener
                .set_nonblocking(true)
                .expect("nonblocking listener");
            let addr = listener.local_addr().expect("local addr");
            let shutdown = Arc::new(AtomicBool::new(false));
            let server_shutdown = Arc::clone(&shutdown);
            let handle = thread::spawn(move || {
                while !server_shutdown.load(Ordering::Relaxed) {
                    match listener.accept() {
                        Ok((mut stream, _)) => {
                            let request = read_stub_request(&mut stream);
                            if request.contains("/embeddings") {
                                write_json_response(
                                    &mut stream,
                                    r#"{"model":"mock-embedding","data":[{"index":0,"embedding":[1,0,0]}],"usage":{"prompt_tokens":3,"total_tokens":3}}"#,
                                );
                            } else {
                                write_json_response(
                                    &mut stream,
                                    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}}"#,
                                );
                            }
                        }
                        Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => {
                            thread::sleep(Duration::from_millis(1));
                        }
                        Err(_) => break,
                    }
                }
            });

            Self {
                addr,
                shutdown,
                handle: Some(handle),
            }
        }

        fn addr(&self) -> SocketAddr {
            self.addr
        }
    }

    impl Drop for AskStub {
        fn drop(&mut self) {
            self.shutdown.store(true, Ordering::Relaxed);
            let _ = TcpStream::connect(self.addr);
            if let Some(handle) = self.handle.take() {
                let _ = handle.join();
            }
        }
    }

    fn read_stub_request(stream: &mut TcpStream) -> String {
        let _ = stream.set_read_timeout(Some(Duration::from_millis(100)));
        let mut buffer = [0_u8; 4096];
        let count = stream.read(&mut buffer).unwrap_or(0);
        String::from_utf8_lossy(&buffer[..count]).into_owned()
    }

    fn write_json_response(stream: &mut TcpStream, body: &str) {
        let response = format!(
            "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
            body.len()
        );
        stream
            .write_all(response.as_bytes())
            .expect("write stub response");
    }
}