leviath-mcp 0.1.0

//! Tool execution via MCP.

use serde_json::Value;
use std::collections::{HashMap, HashSet};

use crate::client::{MCPClient, ToolResult, ToolResultContent};
use crate::discovery::ToolMetadata;

/// Provider tool-name limit: the name advertised to the LLM must match
/// `^[A-Za-z0-9_-]{1,64}$` (the Anthropic/OpenAI rule). MCP names are laxer
/// (they allow dots), so any MCP name that violates this would make the
/// provider reject the *entire* request.
const MAX_TOOL_NAME_LEN: usize = 64;

/// Sanitize an MCP tool name into the provider-accepted character set.
///
/// Every character outside `[A-Za-z0-9_-]` (notably `.`, which MCP allows and
/// real servers use) becomes `_`, and the result is truncated to 64 bytes. An
/// empty result (a name of only illegal characters) falls back to `tool`.
pub fn sanitize_tool_name(name: &str) -> String {
    let mut out: String = name
        .chars()
        .map(|c| {
            if c.is_ascii_alphanumeric() || c == '_' || c == '-' {
                c
            } else {
                '_'
            }
        })
        .collect();
    out.truncate(MAX_TOOL_NAME_LEN);
    if out.is_empty() {
        "tool".to_string()
    } else {
        out
    }
}

/// Result of a tool execution, with convenience fields.
#[derive(Debug, Clone)]
pub struct ExecutionResult {
    /// Whether execution succeeded
    pub success: bool,
    /// Result data as JSON (contains the content array)
    pub data: Value,
    /// Concatenated text content for convenience
    pub text: String,
}

/// Tool execution service that routes tool calls to the correct MCP server.
pub struct ToolExecutor {
    /// Active MCP clients, keyed by server name
    clients: HashMap<String, MCPClient>,
    /// Advertised tool name → (server name, original tool name).
    ///
    /// The name advertised to the LLM is sanitized to the provider's character
    /// rule and made unique across servers; this maps it back to the server and
    /// the original name the server itself expects on a `tools/call`.
    aliases: HashMap<String, (String, String)>,
}

impl ToolExecutor {
    /// Create a new tool executor.
    pub fn new() -> Self {
        Self {
            clients: HashMap::new(),
            aliases: HashMap::new(),
        }
    }

    /// Register an MCP client for a server, advertising its tools under names
    /// safe for the LLM/provider.
    ///
    /// Equivalent to [`Self::add_client_advertised`] reserving nothing; kept for
    /// callers that don't need the advertised metadata back.
    pub fn add_client(&mut self, server_name: String, client: MCPClient) {
        let _ = self.add_client_advertised(server_name, client, &HashSet::new());
    }

    /// Register a client and return its tools under *advertised* names.
    ///
    /// Each advertised name is [`sanitize_tool_name`]d and made unique against
    /// `reserved` (e.g. the built-in tool names) and every previously-registered
    /// tool, so the set of names handed to the provider is always valid and
    /// collision-free. The returned metadata carries the advertised names; the
    /// alias back to `(server, original)` is recorded for routing.
    pub fn add_client_advertised(
        &mut self,
        server_name: String,
        client: MCPClient,
        reserved: &HashSet<String>,
    ) -> Vec<ToolMetadata> {
        let mut advertised = Vec::new();
        for tool in client.cached_tools() {
            let name = self.unique_advertised_name(&tool.name, &server_name, reserved);
            self.aliases
                .insert(name.clone(), (server_name.clone(), tool.name.clone()));
            if name != tool.name {
                tracing::debug!(
                    server = %server_name,
                    original = %tool.name,
                    advertised = %name,
                    "Renamed MCP tool to satisfy provider naming rules"
                );
            }
            advertised.push(ToolMetadata {
                name,
                description: tool.description.clone(),
                schema: tool.schema.clone(),
            });
        }
        self.clients.insert(server_name, client);
        advertised
    }

    /// Compute a unique, provider-safe advertised name for `original`.
    ///
    /// Prefers the sanitized name; on a clash with `reserved` or an existing
    /// alias, prefixes with the server name; if that still clashes, appends a
    /// numeric suffix.
    fn unique_advertised_name(
        &self,
        original: &str,
        server: &str,
        reserved: &HashSet<String>,
    ) -> String {
        let free = |name: &str| !reserved.contains(name) && !self.aliases.contains_key(name);

        let base = sanitize_tool_name(original);
        if free(&base) {
            return base;
        }
        let prefixed = sanitize_tool_name(&format!("{server}__{original}"));
        if free(&prefixed) {
            return prefixed;
        }
        let mut n = 2;
        loop {
            let candidate = sanitize_tool_name(&format!("{prefixed}_{n}"));
            if free(&candidate) {
                return candidate;
            }
            n += 1;
        }
    }

    /// Execute a tool by its advertised name, routing to the owning server.
    pub async fn execute(
        &mut self,
        tool_name: &str,
        arguments: Value,
    ) -> anyhow::Result<ExecutionResult> {
        tracing::info!(tool = %tool_name, "Executing tool");

        // The advertised → (server, original) alias is the authoritative route.
        if let Some((server, original)) = self.aliases.get(tool_name).cloned() {
            return self.execute_on(&server, &original, arguments).await;
        }

        Err(anyhow::anyhow!(
            "No MCP server found with tool '{}'. Available tools: {:?}",
            tool_name,
            self.aliases.keys().collect::<Vec<_>>()
        ))
    }

    /// Execute a tool on a specific server.
    pub async fn execute_on(
        &mut self,
        server_name: &str,
        tool_name: &str,
        arguments: Value,
    ) -> anyhow::Result<ExecutionResult> {
        tracing::info!(server = %server_name, tool = %tool_name, "Executing tool on server");

        let client = self
            .clients
            .get_mut(server_name)
            .ok_or_else(|| anyhow::anyhow!("MCP server '{}' not found", server_name))?;

        let tool_result = client.call_tool(tool_name, arguments).await?;
        Ok(Self::map_result(tool_result))
    }

    /// Execute a tool only if it is in the allowed list.
    ///
    /// Returns an error if the tool is not in the allowed_tools list.
    pub async fn execute_filtered(
        &mut self,
        tool_name: &str,
        arguments: Value,
        allowed_tools: &[String],
    ) -> anyhow::Result<ExecutionResult> {
        if !allowed_tools.iter().any(|t| t == tool_name) {
            return Ok(ExecutionResult {
                success: false,
                data: Value::Null,
                text: format!(
                    "Tool '{}' is not allowed in the current stage. Allowed tools: {:?}",
                    tool_name, allowed_tools
                ),
            });
        }
        self.execute(tool_name, arguments).await
    }

    /// Shutdown all connected MCP clients.
    ///
    /// `MCPClient::shutdown` always returns `Ok` by design (it swallows
    /// subprocess errors so a dead server cannot block cleanup), so errors
    /// are discarded here too.
    pub async fn shutdown_all(&mut self) -> anyhow::Result<()> {
        tracing::info!("Shutting down all MCP clients");
        for client in self.clients.values_mut() {
            let _ = client.shutdown().await;
        }
        self.clients.clear();
        Ok(())
    }

    /// Get the number of connected servers.
    pub fn server_count(&self) -> usize {
        self.clients.len()
    }

    /// Map a ToolResult into an ExecutionResult.
    ///
    /// Only the model-readable blocks contribute to `text`; binary payloads
    /// (image/audio) and bare resource links do not, and an unmodelled block is
    /// skipped with a warning rather than failing the call.
    fn map_result(tool_result: ToolResult) -> ExecutionResult {
        let mut parts: Vec<&str> = Vec::new();
        for content in &tool_result.content {
            match content {
                ToolResultContent::Text { text } => parts.push(text.as_str()),
                ToolResultContent::Resource { resource } => {
                    if let Some(text) = resource.text.as_deref() {
                        parts.push(text);
                    }
                }
                ToolResultContent::Image { .. }
                | ToolResultContent::Audio { .. }
                | ToolResultContent::ResourceLink { .. } => {}
                ToolResultContent::Unknown => {
                    tracing::warn!("Skipping unrecognized MCP content block in tool result");
                }
            }
        }
        let mut text = parts.join("\n");

        // A structured-only result would otherwise reach the model as an empty
        // string. Servers *should* also mirror it into a text block, but not
        // all do.
        if text.is_empty()
            && let Some(structured) = &tool_result.structured_content
        {
            text = structured.to_string();
        }

        let data = serde_json::to_value(&tool_result.content).unwrap_or(Value::Null);

        ExecutionResult {
            success: !tool_result.is_error,
            data,
            text,
        }
    }
}

impl Default for ToolExecutor {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::client::EmbeddedResource;
    use crate::test_support::always_on_tracing_guard;

    #[tokio::test]
    async fn test_execute_filtered_rejects_disallowed_tool() {
        let mut executor = ToolExecutor::new();
        let allowed = vec!["read_file".to_string(), "write_file".to_string()];

        let result = executor
            .execute_filtered("delete_file", serde_json::json!({}), &allowed)
            .await
            .unwrap();

        assert!(!result.success);
        assert!(result.text.contains("not allowed"));
    }

    #[test]
    fn test_tool_executor_creation() {
        let executor = ToolExecutor::new();
        assert_eq!(executor.server_count(), 0);
    }

    // ─── ToolExecutor::default ──────────────────────────────────────────

    #[test]
    fn test_tool_executor_default() {
        let executor = ToolExecutor::default();
        assert_eq!(executor.server_count(), 0);
    }

    // ─── map_result: text content ───────────────────────────────────────

    #[test]
    fn test_map_result_text_content() {
        let tool_result = ToolResult {
            content: vec![ToolResultContent::Text {
                text: "Hello world".to_string(),
            }],
            structured_content: None,
            is_error: false,
        };
        let result = ToolExecutor::map_result(tool_result);
        assert!(result.success);
        assert_eq!(result.text, "Hello world");
    }

    #[test]
    fn test_map_result_error() {
        let tool_result = ToolResult {
            content: vec![ToolResultContent::Text {
                text: "Something failed".to_string(),
            }],
            structured_content: None,
            is_error: true,
        };
        let result = ToolExecutor::map_result(tool_result);
        assert!(!result.success);
        assert_eq!(result.text, "Something failed");
    }

    #[test]
    fn test_map_result_empty_content() {
        let tool_result = ToolResult {
            content: vec![],
            structured_content: None,
            is_error: false,
        };
        let result = ToolExecutor::map_result(tool_result);
        assert!(result.success);
        assert_eq!(result.text, "");
    }

    #[test]
    fn test_map_result_multiple_text() {
        let tool_result = ToolResult {
            content: vec![
                ToolResultContent::Text {
                    text: "line1".to_string(),
                },
                ToolResultContent::Text {
                    text: "line2".to_string(),
                },
            ],
            structured_content: None,
            is_error: false,
        };
        let result = ToolExecutor::map_result(tool_result);
        assert_eq!(result.text, "line1\nline2");
    }

    #[test]
    fn test_map_result_image_excluded_from_text() {
        let tool_result = ToolResult {
            content: vec![
                ToolResultContent::Text {
                    text: "before".to_string(),
                },
                ToolResultContent::Image {
                    data: "base64data".to_string(),
                    mime_type: "image/png".to_string(),
                },
                ToolResultContent::Text {
                    text: "after".to_string(),
                },
            ],
            structured_content: None,
            is_error: false,
        };
        let result = ToolExecutor::map_result(tool_result);
        assert_eq!(result.text, "before\nafter");
    }

    #[test]
    fn test_map_result_resource_with_text() {
        let tool_result = ToolResult {
            content: vec![ToolResultContent::Resource {
                resource: EmbeddedResource {
                    uri: "file:///test".to_string(),
                    text: Some("resource content".to_string()),
                    blob: None,
                    mime_type: None,
                },
            }],
            structured_content: None,
            is_error: false,
        };
        let result = ToolExecutor::map_result(tool_result);
        assert_eq!(result.text, "resource content");
    }

    #[test]
    fn test_map_result_resource_without_text() {
        let tool_result = ToolResult {
            content: vec![ToolResultContent::Resource {
                resource: EmbeddedResource {
                    uri: "file:///test".to_string(),
                    text: None,
                    blob: None,
                    mime_type: None,
                },
            }],
            structured_content: None,
            is_error: false,
        };
        let result = ToolExecutor::map_result(tool_result);
        assert_eq!(result.text, "");
    }

    #[test]
    fn test_map_result_data_is_json() {
        let tool_result = ToolResult {
            content: vec![ToolResultContent::Text {
                text: "hi".to_string(),
            }],
            structured_content: None,
            is_error: false,
        };
        let result = ToolExecutor::map_result(tool_result);
        assert!(result.data.is_array());
    }

    // ─── execute_filtered: allowed tool ────────────────────────────────

    #[tokio::test]
    async fn test_execute_filtered_allowed_tool_but_no_server() {
        let _guard = always_on_tracing_guard();
        let mut executor = ToolExecutor::new();
        let allowed = vec!["read_file".to_string()];

        // Tool is allowed but no server has it
        let result = executor
            .execute_filtered("read_file", serde_json::json!({}), &allowed)
            .await;
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("No MCP server"));
    }

    // ─── execute: no server ─────────────────────────────────────────────

    #[tokio::test]
    async fn test_execute_no_server_errors() {
        let _guard = always_on_tracing_guard();
        let mut executor = ToolExecutor::new();
        let result = executor
            .execute("nonexistent_tool", serde_json::json!({}))
            .await;
        assert!(result.is_err());
    }

    // ─── execute_on: unknown server ─────────────────────────────────────

    #[tokio::test]
    async fn test_execute_on_unknown_server() {
        let _guard = always_on_tracing_guard();
        let mut executor = ToolExecutor::new();
        let result = executor
            .execute_on("unknown_server", "tool", serde_json::json!({}))
            .await;
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("not found"));
    }

    // ─── shutdown_all: empty executor ───────────────────────────────────

    #[tokio::test]
    async fn test_shutdown_all_empty() {
        let _guard = always_on_tracing_guard();
        let mut executor = ToolExecutor::new();
        let result = executor.shutdown_all().await;
        assert!(result.is_ok());
        assert_eq!(executor.server_count(), 0);
    }

    // ─── add_client / execute / execute_on with a live client ───────────
    //
    // Same Python-backed JSON-RPC stub approach used in client.rs/discovery.rs
    // tests. Note: MCPClient::shutdown() always returns Ok(()) by design (it
    // swallows failures so a dead server can't block cleanup) - so
    // shutdown_all()'s error-collection branch is intentionally left
    // uncovered here; there's no way to make client.shutdown() fail without
    // changing that documented "always succeeds" behavior.

    const STUB_INIT_LIST_AND_CALL: &str = r#"
import sys, json

def respond(id, result):
    msg = json.dumps({"jsonrpc": "2.0", "id": id, "result": result})
    sys.stdout.write(msg + "\n")
    sys.stdout.flush()

for line in sys.stdin:
    line = line.strip()
    if not line:
        continue
    req = json.loads(line)
    method = req.get("method", "")
    id_ = req.get("id")
    if method == "initialize":
        respond(id_, {"capabilities": {"tools": {"listChanged": True}}, "protocolVersion": "2024-11-05"})
    elif method == "notifications/initialized":
        pass
    elif method == "tools/list":
        respond(id_, {"tools": [{"name": "echo", "description": "echo tool", "inputSchema": {}}]})
    elif method == "tools/call":
        respond(id_, {"content": [{"type": "text", "text": "hello from tool"}], "isError": False})
    elif method == "notifications/cancelled":
        pass
    else:
        respond(id_, {"error": {"code": -32601, "message": "method not found"}})
"#;

    async fn spawn_ready_client() -> MCPClient {
        let mut client =
            MCPClient::spawn("python3", &["-c", STUB_INIT_LIST_AND_CALL], &HashMap::new())
                .await
                .expect("failed to spawn stub server");
        client.connect().await.expect("connect should succeed");
        client
            .list_tools()
            .await
            .expect("list_tools should succeed");
        client
    }

    #[tokio::test]
    async fn add_client_and_server_count_reflects_it() {
        let mut executor = ToolExecutor::new();
        let client = spawn_ready_client().await;
        executor.add_client("server1".to_string(), client);
        assert_eq!(executor.server_count(), 1);
    }

    #[tokio::test]
    async fn execute_finds_owning_server_and_calls_tool() {
        let _guard = always_on_tracing_guard();
        let mut executor = ToolExecutor::new();
        let client = spawn_ready_client().await;
        executor.add_client("server1".to_string(), client);

        let result = executor
            .execute("echo", serde_json::json!({"text": "hi"}))
            .await
            .expect("execute should succeed");
        assert!(result.success);
        assert_eq!(result.text, "hello from tool");
    }

    #[tokio::test]
    async fn execute_on_specific_server_calls_tool() {
        let _guard = always_on_tracing_guard();
        let mut executor = ToolExecutor::new();
        let client = spawn_ready_client().await;
        executor.add_client("server1".to_string(), client);

        let result = executor
            .execute_on("server1", "echo", serde_json::json!({}))
            .await
            .expect("execute_on should succeed");
        assert!(result.success);
        assert_eq!(result.text, "hello from tool");
    }

    #[tokio::test]
    async fn execute_filtered_allowed_tool_with_server_succeeds() {
        let _guard = always_on_tracing_guard();
        let mut executor = ToolExecutor::new();
        let client = spawn_ready_client().await;
        executor.add_client("server1".to_string(), client);

        let allowed = vec!["echo".to_string()];
        let result = executor
            .execute_filtered("echo", serde_json::json!({}), &allowed)
            .await
            .expect("execute_filtered should succeed");
        assert!(result.success);
    }

    #[tokio::test]
    async fn shutdown_all_with_live_client_succeeds_and_clears() {
        let _guard = always_on_tracing_guard();
        let mut executor = ToolExecutor::new();
        let client = spawn_ready_client().await;
        executor.add_client("server1".to_string(), client);

        let result = executor.shutdown_all().await;
        assert!(result.is_ok());
        assert_eq!(executor.server_count(), 0);
    }

    // ─── ExecutionResult ────────────────────────────────────────────────

    #[test]
    fn test_execution_result_clone() {
        let result = ExecutionResult {
            success: true,
            data: serde_json::json!("test"),
            text: "hello".to_string(),
        };
        let cloned = result.clone();
        assert!(cloned.success);
        assert_eq!(cloned.text, "hello");
    }

    #[test]
    fn test_execution_result_debug() {
        let result = ExecutionResult {
            success: false,
            data: Value::Null,
            text: "error".to_string(),
        };
        let debug = format!("{:?}", result);
        assert!(debug.contains("success"));
        assert!(debug.contains("false"));
    }

    // ─── execute_on: call_tool error propagation ────────────────────────
    //
    // Server returns a JSON-RPC error for tools/call, which causes
    // execute_on's `client.call_tool(...).await?` to propagate the error.

    const STUB_CALL_ERROR: &str = r#"
import sys, json

def respond(id, result):
    msg = json.dumps({"jsonrpc": "2.0", "id": id, "result": result})
    sys.stdout.write(msg + "\n")
    sys.stdout.flush()

def error(id, message):
    msg = json.dumps({"jsonrpc": "2.0", "id": id, "error": {"code": -32603, "message": message}})
    sys.stdout.write(msg + "\n")
    sys.stdout.flush()

for line in sys.stdin:
    line = line.strip()
    if not line:
        continue
    req = json.loads(line)
    method = req.get("method", "")
    id_ = req.get("id")
    if method == "initialize":
        respond(id_, {"capabilities": {"tools": {}}, "protocolVersion": "2024-11-05"})
    elif method == "notifications/initialized":
        pass
    elif method == "tools/list":
        respond(id_, {"tools": [{"name": "echo", "description": "echo", "inputSchema": {}}]})
    elif method == "tools/call":
        error(id_, "tool execution failed")
    elif method == "notifications/cancelled":
        pass
"#;

    #[tokio::test]
    async fn execute_on_propagates_call_tool_error() {
        let _guard = always_on_tracing_guard();
        let mut client = MCPClient::spawn("python3", &["-c", STUB_CALL_ERROR], &HashMap::new())
            .await
            .expect("spawn");
        client.connect().await.expect("connect");
        client.list_tools().await.expect("list_tools");

        let mut executor = ToolExecutor::new();
        executor.add_client("server1".to_string(), client);

        let result = executor
            .execute_on("server1", "echo", serde_json::json!({}))
            .await;
        assert!(result.is_err());
        assert!(
            result
                .unwrap_err()
                .to_string()
                .contains("tool execution failed")
        );
    }

    // ─── map_result over the full content-block set ───────────────────────

    fn text_of(content: Vec<ToolResultContent>) -> String {
        ToolExecutor::map_result(ToolResult {
            content,
            structured_content: None,
            is_error: false,
        })
        .text
    }

    #[test]
    fn map_result_reports_tool_execution_error_as_failure() {
        // The end-to-end consequence of the `isError` rename: a failing tool
        // must reach the model as a failure, not a success.
        let result = ToolExecutor::map_result(ToolResult {
            content: vec![ToolResultContent::Text {
                text: "Invalid departure date".to_string(),
            }],
            structured_content: None,
            is_error: true,
        });
        assert!(!result.success);
        assert_eq!(result.text, "Invalid departure date");
    }

    #[test]
    fn map_result_skips_binary_blocks() {
        let text = text_of(vec![
            ToolResultContent::Text {
                text: "before".to_string(),
            },
            ToolResultContent::Image {
                data: "YWJj".to_string(),
                mime_type: "image/png".to_string(),
            },
            ToolResultContent::Audio {
                data: "YWJj".to_string(),
                mime_type: "audio/wav".to_string(),
            },
            ToolResultContent::Text {
                text: "after".to_string(),
            },
        ]);
        assert_eq!(text, "before\nafter");
    }

    #[test]
    fn map_result_skips_resource_links() {
        let text = text_of(vec![ToolResultContent::ResourceLink {
            uri: "file:///x".to_string(),
            name: "x".to_string(),
            description: None,
            mime_type: None,
        }]);
        assert_eq!(text, "");
    }

    #[test]
    fn map_result_skips_unknown_blocks_without_losing_the_rest() {
        let _guard = always_on_tracing_guard();
        let text = text_of(vec![
            ToolResultContent::Unknown,
            ToolResultContent::Text {
                text: "still here".to_string(),
            },
        ]);
        assert_eq!(text, "still here");
    }

    #[test]
    fn map_result_falls_back_to_structured_content_when_no_text() {
        // Servers *should* mirror structured output into a text block; not all
        // do, and without this the model would receive an empty string.
        let result = ToolExecutor::map_result(ToolResult {
            content: vec![],
            structured_content: Some(serde_json::json!({"temperature": 22.5})),
            is_error: false,
        });
        assert_eq!(result.text, r#"{"temperature":22.5}"#);
    }

    #[test]
    fn map_result_prefers_text_blocks_over_structured_content() {
        let result = ToolExecutor::map_result(ToolResult {
            content: vec![ToolResultContent::Text {
                text: "human readable".to_string(),
            }],
            structured_content: Some(serde_json::json!({"a": 1})),
            is_error: false,
        });
        assert_eq!(result.text, "human readable");
    }

    #[test]
    fn map_result_embedded_resource_blob_contributes_no_text() {
        let text = text_of(vec![ToolResultContent::Resource {
            resource: EmbeddedResource {
                uri: "file:///a.png".to_string(),
                text: None,
                blob: Some("YWJj".to_string()),
                mime_type: Some("image/png".to_string()),
            },
        }]);
        assert_eq!(text, "");
    }

    // ─── tool-name sanitization ───────────────────────────────────────────

    #[test]
    fn sanitize_passes_a_clean_name_through() {
        assert_eq!(sanitize_tool_name("get_weather-2"), "get_weather-2");
    }

    #[test]
    fn sanitize_replaces_dots_and_other_illegal_chars() {
        // Dots are legal in MCP but rejected by the provider name rule.
        assert_eq!(sanitize_tool_name("admin.tools.list"), "admin_tools_list");
        assert_eq!(sanitize_tool_name("weird name!/#"), "weird_name___");
    }

    #[test]
    fn sanitize_truncates_to_the_limit() {
        let long = "a".repeat(200);
        assert_eq!(sanitize_tool_name(&long).len(), MAX_TOOL_NAME_LEN);
    }

    #[test]
    fn sanitize_of_illegal_chars_becomes_underscores_and_empty_falls_back() {
        // Illegal chars each become `_` (still a valid name); only a fully
        // empty result falls back to a placeholder.
        assert_eq!(sanitize_tool_name("...."), "____");
        assert_eq!(sanitize_tool_name(""), "tool");
    }

    // ─── unique_advertised_name ───────────────────────────────────────────

    #[test]
    fn unique_name_prefers_the_sanitized_base() {
        let exec = ToolExecutor::new();
        let reserved = HashSet::new();
        assert_eq!(
            exec.unique_advertised_name("github.search", "gh", &reserved),
            "github_search"
        );
    }

    #[test]
    fn unique_name_prefixes_on_a_reserved_collision() {
        // The base clashes with a built-in tool name → prefix with the server.
        let exec = ToolExecutor::new();
        let reserved: HashSet<String> = ["bash".to_string()].into_iter().collect();
        assert_eq!(
            exec.unique_advertised_name("bash", "srv", &reserved),
            "srv__bash"
        );
    }

    #[test]
    fn unique_name_prefixes_on_an_existing_alias_collision() {
        let mut exec = ToolExecutor::new();
        exec.aliases.insert(
            "search".to_string(),
            ("a".to_string(), "search".to_string()),
        );
        assert_eq!(
            exec.unique_advertised_name("search", "b", &HashSet::new()),
            "b__search"
        );
    }

    #[test]
    fn unique_name_appends_a_number_when_the_prefix_also_collides() {
        let mut exec = ToolExecutor::new();
        // Both the base and the server-prefixed form are already taken.
        exec.aliases.insert(
            "search".to_string(),
            ("a".to_string(), "search".to_string()),
        );
        exec.aliases
            .insert("b__search".to_string(), ("x".to_string(), "y".to_string()));
        assert_eq!(
            exec.unique_advertised_name("search", "b", &HashSet::new()),
            "b__search_2"
        );

        // And when _2 is taken too, it moves on to _3 (covers the loop step).
        exec.aliases.insert(
            "b__search_2".to_string(),
            ("x".to_string(), "y".to_string()),
        );
        assert_eq!(
            exec.unique_advertised_name("search", "b", &HashSet::new()),
            "b__search_3"
        );
    }

    // ─── advertised routing with live clients ─────────────────────────────

    /// A stub whose single tool is named `tool_name`, echoing a fixed reply.
    fn stub_named(tool_name: &str) -> String {
        format!(
            r#"
import sys, json
def respond(id, result):
    sys.stdout.write(json.dumps({{"jsonrpc": "2.0", "id": id, "result": result}}) + "\n")
    sys.stdout.flush()
for line in sys.stdin:
    line = line.strip()
    if not line:
        continue
    req = json.loads(line)
    method, id_ = req.get("method", ""), req.get("id")
    if method == "initialize":
        respond(id_, {{"capabilities": {{}}, "protocolVersion": "2024-11-05"}})
    elif method == "tools/list":
        respond(id_, {{"tools": [{{"name": "{tool_name}", "inputSchema": {{}}}}]}})
    elif method == "tools/call":
        respond(id_, {{"content": [{{"type": "text", "text": "called " + req["params"]["name"]}}], "isError": False}})
"#
        )
    }

    async fn spawn_named(tool_name: &str) -> MCPClient {
        let mut client =
            MCPClient::spawn("python3", &["-c", &stub_named(tool_name)], &HashMap::new())
                .await
                .expect("spawn");
        client.connect().await.expect("connect");
        client.list_tools().await.expect("list");
        client
    }

    #[tokio::test]
    async fn a_dotted_tool_is_advertised_sanitized_and_still_routes() {
        let _guard = always_on_tracing_guard();
        let mut executor = ToolExecutor::new();
        let client = spawn_named("github.search").await;
        let advertised = executor.add_client_advertised("gh".to_string(), client, &HashSet::new());
        assert_eq!(advertised[0].name, "github_search");

        // The LLM calls the advertised name; the server is called with its
        // original name ("github.search").
        let result = executor
            .execute("github_search", serde_json::json!({}))
            .await
            .expect("advertised name routes");
        assert!(result.success);
        assert_eq!(result.text, "called github.search");
    }

    #[tokio::test]
    async fn two_servers_sharing_a_tool_name_are_disambiguated() {
        let _guard = always_on_tracing_guard();
        let mut executor = ToolExecutor::new();

        let a = spawn_named("search").await;
        let a_names = executor.add_client_advertised("alpha".to_string(), a, &HashSet::new());
        assert_eq!(a_names[0].name, "search");

        let b = spawn_named("search").await;
        // Reserve what alpha already advertised.
        let reserved: HashSet<String> = a_names.iter().map(|t| t.name.clone()).collect();
        let b_names = executor.add_client_advertised("beta".to_string(), b, &reserved);
        assert_eq!(b_names[0].name, "beta__search");

        // Both route to their own server with the original name "search".
        assert!(
            executor
                .execute("search", serde_json::json!({}))
                .await
                .unwrap()
                .success
        );
        assert!(
            executor
                .execute("beta__search", serde_json::json!({}))
                .await
                .unwrap()
                .success
        );
    }

    #[tokio::test]
    async fn add_client_reserving_nothing_registers_identity_aliases() {
        let _guard = always_on_tracing_guard();
        let mut executor = ToolExecutor::new();
        let client = spawn_named("plain").await;
        executor.add_client("s".to_string(), client);
        assert!(
            executor
                .execute("plain", serde_json::json!({}))
                .await
                .unwrap()
                .success
        );
    }
}