161 lines · 4.8 KB · Rust
RawDownload
pub mod client;pub mod listener;pub mod protocol;#[cfg(any(test, feature = "test-support"))]pub mod test;pub mod transport;pub mod types; use collections::HashMap;use http_client::HttpClient;use std::path::Path;use std::sync::Arc;use std::time::Duration;use std::{fmt::Display, path::PathBuf}; use anyhow::Result;use client::Client;use gpui::AsyncApp;use parking_lot::RwLock;pub use settings::ContextServerCommand;use url::Url; use crate::transport::HttpTransport; #[derive(Debug, Clone, PartialEq, Eq, Hash)]pub struct ContextServerId(pub Arc<str>); impl Display for ContextServerId {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        write!(f, "{}", self.0)    }} enum ContextServerTransport {    Stdio(ContextServerCommand, Option<PathBuf>),    Custom(Arc<dyn crate::transport::Transport>),} pub struct ContextServer {    id: ContextServerId,    client: RwLock<Option<Arc<crate::protocol::InitializedContextServerProtocol>>>,    configuration: ContextServerTransport,    request_timeout: Option<Duration>,} impl ContextServer {    pub fn stdio(        id: ContextServerId,        command: ContextServerCommand,        working_directory: Option<Arc<Path>>,    ) -> Self {        Self {            id,            client: RwLock::new(None),            configuration: ContextServerTransport::Stdio(                command,                working_directory.map(|directory| directory.to_path_buf()),            ),            request_timeout: None,        }    }     pub fn http(        id: ContextServerId,        endpoint: &Url,        headers: HashMap<String, String>,        http_client: Arc<dyn HttpClient>,        executor: gpui::BackgroundExecutor,        request_timeout: Option<Duration>,    ) -> Result<Self> {        let transport = match endpoint.scheme() {            "http" | "https" => {                log::info!("Using HTTP transport for {}", endpoint);                let transport =                    HttpTransport::new(http_client, endpoint.to_string(), headers, executor);                Arc::new(transport) as _            }            _ => anyhow::bail!("unsupported MCP url scheme {}", endpoint.scheme()),        };        Ok(Self::new_with_timeout(id, transport, request_timeout))    }     pub fn new(id: ContextServerId, transport: Arc<dyn crate::transport::Transport>) -> Self {        Self::new_with_timeout(id, transport, None)    }     pub fn new_with_timeout(        id: ContextServerId,        transport: Arc<dyn crate::transport::Transport>,        request_timeout: Option<Duration>,    ) -> Self {        Self {            id,            client: RwLock::new(None),            configuration: ContextServerTransport::Custom(transport),            request_timeout,        }    }     pub fn id(&self) -> ContextServerId {        self.id.clone()    }     pub fn client(&self) -> Option<Arc<crate::protocol::InitializedContextServerProtocol>> {        self.client.read().clone()    }     pub async fn start(&self, cx: &AsyncApp) -> Result<()> {        self.initialize(self.new_client(cx)?).await    }     fn new_client(&self, cx: &AsyncApp) -> Result<Client> {        Ok(match &self.configuration {            ContextServerTransport::Stdio(command, working_directory) => Client::stdio(                client::ContextServerId(self.id.0.clone()),                client::ModelContextServerBinary {                    executable: Path::new(&command.path).to_path_buf(),                    args: command.args.clone(),                    env: command.env.clone(),                    timeout: command.timeout,                },                working_directory,                cx.clone(),            )?,            ContextServerTransport::Custom(transport) => Client::new(                client::ContextServerId(self.id.0.clone()),                self.id().0,                transport.clone(),                self.request_timeout,                cx.clone(),            )?,        })    }     async fn initialize(&self, client: Client) -> Result<()> {        log::debug!("starting context server {}", self.id);        let protocol = crate::protocol::ModelContextProtocol::new(client);        let client_info = types::Implementation {            name: "Zed".to_string(),            version: env!("CARGO_PKG_VERSION").to_string(),        };        let initialized_protocol = protocol.initialize(client_info).await?;         log::debug!(            "context server {} initialized: {:?}",            self.id,            initialized_protocol.initialize,        );         *self.client.write() = Some(Arc::new(initialized_protocol));        Ok(())    }     pub fn stop(&self) -> Result<()> {        let mut client = self.client.write();        if let Some(protocol) = client.take() {            drop(protocol);        }        Ok(())    }}