tenant.openagents/omega

1pub mod client;
2pub mod listener;
3pub mod oauth;
4pub mod protocol;
5#[cfg(any(test, feature = "test-support"))]
6pub mod test;
7pub mod transport;
8pub mod types;
9
10use collections::HashMap;
11use http_client::HttpClient;
12use std::path::Path;
13use std::sync::Arc;
14use std::time::Duration;
15use std::{fmt::Display, path::PathBuf};
16
17use anyhow::Result;
18use client::Client;
19use gpui::AsyncApp;
20use parking_lot::RwLock;
21pub use settings::ContextServerCommand;
22use url::Url;
23
24use crate::oauth::WwwAuthenticate;
25use crate::transport::HttpTransport;
26
27#[derive(Debug, Clone, PartialEq, Eq, Hash)]
28pub struct ContextServerId(pub Arc<str>);
29
30impl Display for ContextServerId {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        write!(f, "{}", self.0)
33    }
34}
35
36enum ContextServerTransport {
37    Stdio(ContextServerCommand, Option<PathBuf>),
38    Custom(Arc<dyn crate::transport::Transport>),
39}
40
41pub struct ContextServer {
42    id: ContextServerId,
43    client: RwLock<Option<Arc<crate::protocol::InitializedContextServerProtocol>>>,
44    configuration: ContextServerTransport,
45    request_timeout: Option<Duration>,
46}
47
48impl ContextServer {
49    pub fn stdio(
50        id: ContextServerId,
51        command: ContextServerCommand,
52        working_directory: Option<Arc<Path>>,
53    ) -> Self {
54        Self {
55            id,
56            client: RwLock::new(None),
57            configuration: ContextServerTransport::Stdio(
58                command,
59                working_directory.map(|directory| directory.to_path_buf()),
60            ),
61            request_timeout: None,
62        }
63    }
64
65    pub fn http(
66        id: ContextServerId,
67        endpoint: &Url,
68        headers: HashMap<String, String>,
69        http_client: Arc<dyn HttpClient>,
70        executor: gpui::BackgroundExecutor,
71        request_timeout: Option<Duration>,
72    ) -> Result<Self> {
73        let transport = match endpoint.scheme() {
74            "http" | "https" => {
75                log::info!("Using HTTP transport for {}", endpoint);
76                let transport =
77                    HttpTransport::new(http_client, endpoint.to_string(), headers, executor);
78                Arc::new(transport) as _
79            }
80            _ => anyhow::bail!("unsupported MCP url scheme {}", endpoint.scheme()),
81        };
82        Ok(Self::new_with_timeout(id, transport, request_timeout))
83    }
84
85    pub fn new(id: ContextServerId, transport: Arc<dyn crate::transport::Transport>) -> Self {
86        Self::new_with_timeout(id, transport, None)
87    }
88
89    pub fn new_with_timeout(
90        id: ContextServerId,
91        transport: Arc<dyn crate::transport::Transport>,
92        request_timeout: Option<Duration>,
93    ) -> Self {
94        Self {
95            id,
96            client: RwLock::new(None),
97            configuration: ContextServerTransport::Custom(transport),
98            request_timeout,
99        }
100    }
101
102    pub fn id(&self) -> ContextServerId {
103        self.id.clone()
104    }
105
106    pub fn client(&self) -> Option<Arc<crate::protocol::InitializedContextServerProtocol>> {
107        self.client.read().clone()
108    }
109
110    /// The authentication challenge from the last `401 Unauthorized` response
111    /// this server's transport gave up on, if any. See
112    /// [`crate::transport::Transport::auth_challenge`].
113    pub fn auth_challenge(&self) -> Option<WwwAuthenticate> {
114        match &self.configuration {
115            ContextServerTransport::Stdio(..) => None,
116            ContextServerTransport::Custom(transport) => transport.auth_challenge(),
117        }
118    }
119
120    pub async fn start(&self, cx: &AsyncApp) -> Result<()> {
121        self.initialize(self.new_client(cx)?).await
122    }
123
124    fn new_client(&self, cx: &AsyncApp) -> Result<Client> {
125        Ok(match &self.configuration {
126            ContextServerTransport::Stdio(command, working_directory) => Client::stdio(
127                client::ContextServerId(self.id.0.clone()),
128                client::ModelContextServerBinary {
129                    executable: Path::new(&command.path).to_path_buf(),
130                    args: command.args.clone(),
131                    env: command.env.clone(),
132                    timeout: command.timeout,
133                },
134                working_directory,
135                cx.clone(),
136            )?,
137            ContextServerTransport::Custom(transport) => Client::new(
138                client::ContextServerId(self.id.0.clone()),
139                self.id().0,
140                transport.clone(),
141                self.request_timeout,
142                cx.clone(),
143            )?,
144        })
145    }
146
147    async fn initialize(&self, client: Client) -> Result<()> {
148        log::debug!("starting context server {}", self.id);
149        let protocol = crate::protocol::ModelContextProtocol::new(client);
150        let client_info = types::Implementation {
151            name: "Zed".to_string(),
152            title: None,
153            version: env!("CARGO_PKG_VERSION").to_string(),
154            description: None,
155        };
156        let initialized_protocol = protocol.initialize(client_info).await?;
157
158        log::debug!(
159            "context server {} initialized: {:?}",
160            self.id,
161            initialized_protocol.initialize,
162        );
163
164        *self.client.write() = Some(Arc::new(initialized_protocol));
165        Ok(())
166    }
167
168    pub fn stop(&self) -> Result<()> {
169        let mut client = self.client.write();
170        if let Some(protocol) = client.take() {
171            drop(protocol);
172        }
173        Ok(())
174    }
175}
176