aether-mcp-utils 0.2.2

use llm::ToolDefinition;

use super::{
    McpError, Result,
    config::{McpServerConfig, ServerConfig},
    connect_mcp::{ConnectOutcome, ConnectionSpec, ProxySpec, Registration, build_plan, connect_mcp},
    connection::{McpServerConnection, ServerInstructions, Tool},
    mcp_client::McpClient,
    naming::{create_namespaced_tool_name, split_on_server_name},
    oauth::{OAuthHandler, perform_oauth_flow},
    tool_proxy::ToolProxy,
};
use futures::future::join_all;
use rmcp::{
    RoleClient,
    model::{
        CallToolRequestParams, ClientCapabilities, ClientInfo, CreateElicitationRequestParams, CreateElicitationResult,
        ElicitationAction, FormElicitationCapability, Implementation, Root, Tool as RmcpTool, UrlElicitationCapability,
    },
    service::RunningService,
    transport::streamable_http_client::StreamableHttpClientTransportConfig,
};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use tokio::sync::{RwLock, mpsc, oneshot};

pub use crate::status::{McpServerStatus, McpServerStatusEntry};

#[derive(Debug)]
pub struct ElicitationRequest {
    pub server_name: String,
    pub request: CreateElicitationRequestParams,
    pub response_sender: oneshot::Sender<CreateElicitationResult>,
}

#[derive(Debug, Clone)]
pub struct ElicitationResponse {
    pub action: ElicitationAction,
    pub content: Option<Value>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct UrlElicitationCompleteParams {
    pub server_name: String,
    pub elicitation_id: String,
}

/// Events emitted by MCP clients that require attention from the host
/// (e.g. the relay or TUI). Flows through a single channel from `McpManager`
/// to the consumer.
#[derive(Debug)]
pub enum McpClientEvent {
    Elicitation(ElicitationRequest),
    UrlElicitationComplete(UrlElicitationCompleteParams),
}

/// Manages connections to multiple MCP servers and their tools
pub struct McpManager {
    servers: HashMap<String, McpServerConnection>,
    tools: HashMap<String, Tool>,
    tool_definitions: Vec<ToolDefinition>,
    client_info: ClientInfo,
    event_sender: mpsc::Sender<McpClientEvent>,
    /// Roots shared with all MCP clients
    roots: Arc<RwLock<Vec<Root>>>,
    oauth_handler: Option<Arc<dyn OAuthHandler>>,
    server_statuses: Vec<McpServerStatusEntry>,
    /// Configs for failed HTTP servers so we can retry OAuth later
    pending_configs: HashMap<String, StreamableHttpClientTransportConfig>,
    /// Optional tool-proxy that wraps multiple servers behind a single `call_tool`.
    proxy: Option<ToolProxy>,
}

impl McpManager {
    pub fn new(event_sender: mpsc::Sender<McpClientEvent>, oauth_handler: Option<Arc<dyn OAuthHandler>>) -> Self {
        let mut capabilities = ClientCapabilities::builder().enable_elicitation().enable_roots().build();
        if let Some(elicitation) = capabilities.elicitation.as_mut() {
            elicitation.form = Some(FormElicitationCapability::default());
            elicitation.url = Some(UrlElicitationCapability::default());
        }

        Self {
            servers: HashMap::new(),
            tools: HashMap::new(),
            tool_definitions: Vec::new(),
            client_info: ClientInfo::new(capabilities, Implementation::new("aether", "0.1.0")),
            event_sender,
            roots: Arc::new(RwLock::new(Vec::new())),
            oauth_handler,
            server_statuses: Vec::new(),
            pending_configs: HashMap::new(),
            proxy: None,
        }
    }

    /// Update or insert the status entry for a server.
    fn set_status(&mut self, name: &str, status: McpServerStatus) {
        if let Some(entry) = self.server_statuses.iter_mut().find(|s| s.name == name) {
            entry.status = status;
        } else {
            self.server_statuses.push(McpServerStatusEntry { name: name.to_string(), status });
        }
    }

    pub async fn add_mcps(&mut self, configs: Vec<McpServerConfig>) -> Result<()> {
        let (direct, proxies) = build_plan(configs).await?;
        let outcomes: Vec<ConnectOutcome> = join_all(direct.into_iter().map(|leaf| {
            connect_mcp(leaf, &self.client_info, &self.event_sender, &self.roots, self.oauth_handler.as_ref())
        }))
        .await;

        let mut mcp_proxies: HashMap<String, HashSet<String>> =
            proxies.iter().map(|p| (p.name.clone(), HashSet::new())).collect();

        let mut connected_mcps_to_proxy: HashMap<String, Vec<String>> = HashMap::new();
        for outcome in outcomes {
            match outcome {
                ConnectOutcome::Ready { name, conn, tools, proxy, registration } => {
                    self.apply_connected(&name, conn, &tools, registration);
                    if let Some(p) = proxy {
                        if let Some(members) = mcp_proxies.get_mut(&p) {
                            members.insert(name.clone());
                        }
                        connected_mcps_to_proxy.entry(p).or_default().push(name);
                    }
                }
                ConnectOutcome::NeedsOAuth { name, config, error, proxy } => {
                    tracing::warn!("Server '{name}' needs OAuth: {error}");
                    self.pending_configs.insert(name.clone(), config);
                    self.set_status(&name, McpServerStatus::NeedsOAuth);
                    if let Some(p) = proxy
                        && let Some(members) = mcp_proxies.get_mut(&p)
                    {
                        members.insert(name);
                    }
                }
                ConnectOutcome::Failed { name, error } => {
                    tracing::warn!("Failed to connect to MCP server '{name}': {error}");
                    if !self.server_statuses.iter().any(|s| s.name == name) {
                        self.set_status(&name, McpServerStatus::Failed { error: error.to_string() });
                    }
                }
            }
        }

        let writes = proxies.iter().flat_map(|proxy| {
            connected_mcps_to_proxy
                .get(&proxy.name)
                .into_iter()
                .flatten()
                .filter_map(|member| {
                    let client = self.servers.get(member)?.client.clone();
                    let dir = proxy.tool_dir.clone();
                    let name = member.clone();
                    Some(async move {
                        if let Err(e) = ToolProxy::write_tools_to_dir(&name, &client, &dir).await {
                            tracing::warn!("Failed to write tool files for nested server '{name}': {e}");
                        }
                    })
                })
                .collect::<Vec<_>>()
        });
        join_all(writes).await;

        for proxy in proxies {
            self.register_proxy(proxy, &mut mcp_proxies, &mut connected_mcps_to_proxy);
        }

        Ok(())
    }

    pub async fn add_mcp_with_auth(&mut self, name: String, base_url: &str, auth_header: String) -> Result<()> {
        let config = ServerConfig::Http {
            name: name.clone(),
            config: StreamableHttpClientTransportConfig::with_uri(base_url).auth_header(auth_header),
        };
        let leaf = ConnectionSpec { name, config, proxy: None, registration: Registration::Direct };
        match connect_mcp(leaf, &self.client_info, &self.event_sender, &self.roots, self.oauth_handler.as_ref()).await {
            ConnectOutcome::Ready { name, conn, tools, registration, .. } => {
                self.apply_connected(&name, conn, &tools, registration);
                Ok(())
            }
            ConnectOutcome::NeedsOAuth { error, .. } | ConnectOutcome::Failed { error, .. } => Err(error),
        }
    }

    pub async fn add_mcp(&mut self, config: McpServerConfig) -> Result<()> {
        match config {
            McpServerConfig::ToolProxy { .. } => self.add_mcps(vec![config]).await,

            McpServerConfig::Server(config) => {
                let leaf = ConnectionSpec {
                    name: config.name().to_string(),
                    config,
                    proxy: None,
                    registration: Registration::Direct,
                };
                match connect_mcp(leaf, &self.client_info, &self.event_sender, &self.roots, self.oauth_handler.as_ref())
                    .await
                {
                    ConnectOutcome::Ready { name, conn, tools, registration, .. } => {
                        self.apply_connected(&name, conn, &tools, registration);
                        Ok(())
                    }
                    ConnectOutcome::NeedsOAuth { name, config, error, .. } => {
                        self.pending_configs.insert(name.clone(), config);
                        self.set_status(&name, McpServerStatus::NeedsOAuth);
                        Err(error)
                    }
                    ConnectOutcome::Failed { error, .. } => Err(error),
                }
            }
        }
    }

    /// Register a proxy's virtual `call_tool`, status, and `self.proxy` slot.
    /// Called after tool files have already been written to disk in a single
    /// cross-proxy parallel batch.
    fn register_proxy(
        &mut self,
        proxy: ProxySpec,
        proxy_members: &mut HashMap<String, HashSet<String>>,
        ready_for_proxy: &mut HashMap<String, Vec<String>>,
    ) {
        let members = ready_for_proxy.remove(&proxy.name).unwrap_or_default();

        let server_descriptions: Vec<(String, String)> = members
            .iter()
            .filter_map(|member| {
                self.servers
                    .get(member)
                    .map(|conn| (member.clone(), ToolProxy::extract_server_description(&conn.client, member)))
            })
            .collect();

        let call_tool_def = ToolProxy::call_tool_definition(&proxy.name);
        self.tools.insert(
            call_tool_def.name.clone(),
            Tool {
                description: call_tool_def.description.clone(),
                parameters: serde_json::from_str(&call_tool_def.parameters)
                    .unwrap_or(Value::Object(serde_json::Map::default())),
            },
        );
        self.tool_definitions.push(call_tool_def);

        let nested = proxy_members.remove(&proxy.name).unwrap_or_default();
        self.proxy = Some(ToolProxy::new(proxy.name.clone(), nested, proxy.tool_dir, &server_descriptions));
        self.set_status(&proxy.name, McpServerStatus::Connected { tool_count: 1 });
    }

    async fn oauth_and_reconnect(&mut self, name: String, config: StreamableHttpClientTransportConfig) -> Result<()> {
        let handler = self
            .oauth_handler
            .as_ref()
            .ok_or_else(|| McpError::ConnectionFailed(format!("No OAuth handler available for '{name}'")))?;
        let auth_client = perform_oauth_flow(&name, &config.uri, handler.as_ref())
            .await
            .map_err(|e| McpError::ConnectionFailed(format!("OAuth failed for '{name}': {e}")))?;

        let mcp_client =
            McpClient::new(self.client_info.clone(), name.clone(), self.event_sender.clone(), Arc::clone(&self.roots));
        let conn = McpServerConnection::reconnect_with_auth(&name, config, auth_client, mcp_client).await?;

        // If this server is proxied, register without exposing tools to the agent
        if let Some(proxy) = self.proxy.as_ref().filter(|p| p.contains_server(&name)) {
            let tool_dir = proxy.tool_dir().to_path_buf();
            self.register_server(&name, conn, Registration::Proxied).await?;
            // Write tool files now that connection succeeded
            if let Some(conn) = self.servers.get(&name) {
                let client = conn.client.clone();
                if let Err(e) = ToolProxy::write_tools_to_dir(&name, &client, &tool_dir).await {
                    tracing::warn!("Failed to write tool files for '{name}' after OAuth: {e}");
                }
            }
            Ok(())
        } else {
            self.register_server(&name, conn, Registration::Direct).await
        }
    }

    /// Register a connected server, discovering its tools first.
    ///
    /// Used by single-server paths (`add_mcp`, `add_mcp_with_auth`,
    /// `oauth_and_reconnect`) where `list_tools` hasn't already been called.
    async fn register_server(
        &mut self,
        name: &str,
        conn: McpServerConnection,
        registration: Registration,
    ) -> Result<()> {
        let tools = conn
            .list_tools()
            .await
            .map_err(|e| McpError::ToolDiscoveryFailed(format!("Failed to list tools for {name}: {e}")))?;
        self.apply_connected(name, conn, &tools, registration);
        Ok(())
    }

    /// Record a connected server whose tools have already been listed.
    ///
    /// When `registration` is `Direct`, discovered tools are added to
    /// `self.tool_definitions` (visible to the agent). When `Proxied`, tools are
    /// only stored in `self.tools` for internal routing.
    fn apply_connected(
        &mut self,
        name: &str,
        conn: McpServerConnection,
        tools: &[RmcpTool],
        registration: Registration,
    ) {
        for rmcp_tool in tools {
            let tool_name = rmcp_tool.name.to_string();
            let namespaced_tool_name = create_namespaced_tool_name(name, &tool_name);
            let tool = Tool::from(rmcp_tool);

            if registration == Registration::Direct {
                self.tool_definitions.push(ToolDefinition {
                    name: namespaced_tool_name.clone(),
                    description: tool.description.clone(),
                    parameters: tool.parameters.to_string(),
                    server: Some(name.to_string()),
                });
            }

            self.tools.insert(namespaced_tool_name, tool);
        }

        self.set_status(name, McpServerStatus::Connected { tool_count: tools.len() });
        self.pending_configs.remove(name);
        self.servers.insert(name.to_string(), conn);
    }

    /// Resolve and route a tool call.
    ///
    /// Returns the target MCP client and normalized call params. For proxy
    /// `call_tool`, this parses the wrapper arguments and forwards to the
    /// selected nested server/tool.
    pub fn get_client_for_tool(
        &self,
        namespaced_tool_name: &str,
        arguments_json: &str,
    ) -> Result<(Arc<RunningService<RoleClient, McpClient>>, CallToolRequestParams)> {
        if !self.tools.contains_key(namespaced_tool_name) {
            return Err(McpError::ToolNotFound(namespaced_tool_name.to_string()));
        }

        let (server_name, tool_name) = split_on_server_name(namespaced_tool_name)
            .ok_or_else(|| McpError::InvalidToolNameFormat(namespaced_tool_name.to_string()))?;

        if let Some(proxy) = self.proxy.as_ref().filter(|p| p.name() == server_name) {
            let call = proxy.resolve_call(arguments_json)?;
            let conn = self
                .servers
                .get(&call.server)
                .ok_or_else(|| McpError::ServerNotFound(format!("Nested server '{}' is not connected", call.server)))?;
            let params = CallToolRequestParams::new(call.tool).with_arguments(call.arguments.unwrap_or_default());
            return Ok((conn.client.clone(), params));
        }

        let client = self
            .servers
            .get(server_name)
            .map(|server| server.client.clone())
            .ok_or_else(|| McpError::ServerNotFound(server_name.to_string()))?;

        let arguments = serde_json::from_str::<serde_json::Value>(arguments_json)?.as_object().cloned();
        let mut params = CallToolRequestParams::new(tool_name.to_string());
        if let Some(args) = arguments {
            params = params.with_arguments(args);
        }

        Ok((client, params))
    }

    pub fn tool_definitions(&self) -> Vec<ToolDefinition> {
        self.tool_definitions.clone()
    }

    /// Returns instructions from all connected MCP servers that provide them,
    /// plus synthesized instructions for tool-proxy groups.
    pub fn server_instructions(&self) -> Vec<ServerInstructions> {
        let mut instructions: Vec<ServerInstructions> = self
            .servers
            .iter()
            .filter(|(name, _)| self.proxy.as_ref().is_none_or(|p| !p.contains_server(name)))
            .filter_map(|(name, conn)| {
                conn.instructions
                    .as_ref()
                    .map(|instr| ServerInstructions { server_name: name.clone(), instructions: instr.clone() })
            })
            .collect();

        if let Some(proxy) = &self.proxy {
            instructions.push(ServerInstructions {
                server_name: proxy.name().to_string(),
                instructions: proxy.instructions().to_string(),
            });
        }

        instructions
    }

    pub fn server_statuses(&self) -> &[McpServerStatusEntry] {
        &self.server_statuses
    }

    /// Authenticate a server that previously failed with `NeedsOAuth`.
    ///
    /// Looks up the pending config, runs the OAuth flow, and updates the status
    /// entry on success.
    pub async fn authenticate_server(&mut self, name: &str) -> Result<()> {
        let config = self
            .pending_configs
            .get(name)
            .ok_or_else(|| McpError::ConnectionFailed(format!("no pending config for server '{name}'")))?
            .clone();

        self.oauth_and_reconnect(name.to_string(), config).await
    }

    /// List all prompts from all connected MCP servers with namespacing
    pub async fn list_prompts(&self) -> Result<Vec<rmcp::model::Prompt>> {
        use futures::future::join_all;

        let futures: Vec<_> = self
            .servers
            .iter()
            .filter(|(_, server_conn)| {
                server_conn.client.peer_info().and_then(|info| info.capabilities.prompts.as_ref()).is_some()
            })
            .map(|(server_name, server_conn)| {
                let server_name = server_name.clone();
                let client = server_conn.client.clone();
                async move {
                    let prompts_response = client.list_prompts(None).await.map_err(|e| {
                        McpError::PromptListFailed(format!("Failed to list prompts for {server_name}: {e}"))
                    })?;

                    let namespaced_prompts: Vec<rmcp::model::Prompt> = prompts_response
                        .prompts
                        .into_iter()
                        .map(|prompt| {
                            let namespaced_name = create_namespaced_tool_name(&server_name, &prompt.name);
                            rmcp::model::Prompt::new(namespaced_name, prompt.description, prompt.arguments)
                        })
                        .collect();

                    Ok::<_, McpError>(namespaced_prompts)
                }
            })
            .collect();

        let results = join_all(futures).await;
        let mut all_prompts = Vec::new();
        for result in results {
            all_prompts.extend(result?);
        }

        Ok(all_prompts)
    }

    /// Get a specific prompt by namespaced name
    pub async fn get_prompt(
        &self,
        namespaced_prompt_name: &str,
        arguments: Option<serde_json::Map<String, serde_json::Value>>,
    ) -> Result<rmcp::model::GetPromptResult> {
        let (server_name, prompt_name) = split_on_server_name(namespaced_prompt_name)
            .ok_or_else(|| McpError::InvalidToolNameFormat(namespaced_prompt_name.to_string()))?;

        let server_conn =
            self.servers.get(server_name).ok_or_else(|| McpError::ServerNotFound(server_name.to_string()))?;

        let mut request = rmcp::model::GetPromptRequestParams::new(prompt_name);
        if let Some(args) = arguments {
            request = request.with_arguments(args);
        }

        server_conn.client.get_prompt(request).await.map_err(|e| {
            McpError::PromptGetFailed(format!("Failed to get prompt '{prompt_name}' from {server_name}: {e}"))
        })
    }

    /// Shutdown all servers and wait for their tasks to complete
    pub async fn shutdown(&mut self) {
        let servers: Vec<(String, McpServerConnection)> = self.servers.drain().collect();

        for (server_name, server) in servers {
            if let Some(handle) = server.server_task {
                // Drop the client first to signal shutdown
                drop(server.client);

                // Wait for the server task to complete (with a timeout)
                match tokio::time::timeout(std::time::Duration::from_secs(5), handle).await {
                    Ok(Ok(())) => {
                        tracing::info!("Server '{server_name}' shut down gracefully");
                    }
                    Ok(Err(e)) => {
                        tracing::warn!("Server '{server_name}' task panicked: {e:?}");
                    }
                    Err(_) => {
                        tracing::warn!("Server '{server_name}' shutdown timed out");
                        // Task will be cancelled when the handle is dropped
                    }
                }
            }
        }

        self.tools.clear();
        self.tool_definitions.clear();
        self.proxy = None;
    }

    /// Shutdown a specific server by name
    pub async fn shutdown_server(&mut self, server_name: &str) -> Result<()> {
        let server = self.servers.remove(server_name);

        if let Some(server) = server {
            if let Some(handle) = server.server_task {
                // Drop the client first to signal shutdown
                drop(server.client);

                // Wait for the server task to complete (with a timeout)
                match tokio::time::timeout(std::time::Duration::from_secs(5), handle).await {
                    Ok(Ok(())) => {
                        tracing::info!("Server '{server_name}' shut down gracefully");
                    }
                    Ok(Err(e)) => {
                        tracing::warn!("Server '{server_name}' task panicked: {e:?}");
                    }
                    Err(_) => {
                        tracing::warn!("Server '{server_name}' shutdown timed out");
                        // Task will be cancelled when the handle is dropped
                    }
                }
            }

            // Remove tools from this server
            self.tools.retain(|tool_name, _| !tool_name.starts_with(server_name));

            self.tool_definitions.retain(|tool_def| !tool_def.name.starts_with(server_name));
        }

        Ok(())
    }

    /// Set the roots advertised to MCP servers.
    ///
    /// This updates the roots and sends notifications to all connected servers
    /// that support the `roots/list_changed` notification.
    pub async fn set_roots(&mut self, new_roots: Vec<Root>) -> Result<()> {
        // Update stored roots
        {
            let mut roots = self.roots.write().await;
            *roots = new_roots;
        }

        // Notify all connected servers
        self.notify_roots_changed().await;

        Ok(())
    }

    /// Send `roots/list_changed` notification to all connected servers.
    ///
    /// This prompts servers to re-request the roots via the roots/list endpoint.
    /// Servers that don't support roots will simply ignore the notification.
    async fn notify_roots_changed(&self) {
        for (server_name, server_conn) in &self.servers {
            // Try to send notification - servers that don't support roots will ignore it
            if let Err(e) = server_conn.client.notify_roots_list_changed().await {
                // Only log errors for debugging; it's expected that some servers may not support roots
                tracing::debug!("Note: server '{server_name}' did not accept roots notification: {e}");
            }
        }
    }
}

impl Drop for McpManager {
    fn drop(&mut self) {
        let servers: Vec<(String, McpServerConnection)> = self.servers.drain().collect();
        for (server_name, server) in servers {
            if let Some(handle) = server.server_task {
                handle.abort();
                tracing::warn!("Server '{server_name}' task aborted during cleanup");
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::McpManager;
    use crate::client::config::ServerConfig;
    use rmcp::{
        Json, RoleServer, ServerHandler,
        handler::server::{router::tool::ToolRouter, wrapper::Parameters},
        model::{Implementation, ServerCapabilities, ServerInfo},
        service::DynService,
        tool, tool_handler, tool_router,
    };
    use schemars::JsonSchema;
    use serde::{Deserialize, Serialize};
    use std::{
        io,
        sync::{Arc, Mutex},
    };
    use tokio::sync::mpsc;

    #[derive(Clone)]
    struct TestServer {
        tool_router: ToolRouter<Self>,
    }

    #[tool_handler(router = self.tool_router)]
    impl ServerHandler for TestServer {
        fn get_info(&self) -> ServerInfo {
            ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
                .with_server_info(Implementation::new("test-server", "0.1.0").with_description("Test MCP server"))
        }
    }

    impl Default for TestServer {
        fn default() -> Self {
            Self { tool_router: Self::tool_router() }
        }
    }

    #[derive(Debug, Deserialize, Serialize, JsonSchema)]
    struct EchoRequest {
        value: String,
    }

    #[derive(Debug, Deserialize, Serialize, JsonSchema)]
    struct EchoResult {
        value: String,
    }

    #[tool_router]
    impl TestServer {
        fn into_dyn(self) -> Box<dyn DynService<RoleServer>> {
            Box::new(self)
        }

        #[tool(description = "Returns the provided value")]
        async fn echo(&self, request: Parameters<EchoRequest>) -> Json<EchoResult> {
            let Parameters(EchoRequest { value }) = request;
            Json(EchoResult { value })
        }
    }

    #[derive(Clone)]
    struct SharedWriter(Arc<Mutex<Vec<u8>>>);

    impl io::Write for SharedWriter {
        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
            self.0.lock().unwrap().extend_from_slice(buf);
            Ok(buf.len())
        }

        fn flush(&mut self) -> io::Result<()> {
            Ok(())
        }
    }

    #[tokio::test]
    async fn drop_logs_cleanup_abort_with_tracing() {
        let (event_sender, _event_receiver) = mpsc::channel(1);
        let mut manager = McpManager::new(event_sender, None);
        manager
            .add_mcp(
                ServerConfig::InMemory { name: "test".to_string(), server: TestServer::default().into_dyn() }.into(),
            )
            .await
            .unwrap();

        let output = Arc::new(Mutex::new(Vec::new()));
        let subscriber = tracing_subscriber::fmt()
            .with_ansi(false)
            .without_time()
            .with_writer({
                let output = Arc::clone(&output);
                move || SharedWriter(Arc::clone(&output))
            })
            .finish();

        tracing::subscriber::with_default(subscriber, || {
            drop(manager);
        });

        let logs = String::from_utf8(output.lock().unwrap().clone()).unwrap();
        assert!(logs.contains("Server 'test' task aborted during cleanup"));
    }
}