1pub mod client;
2pub mod protocol;
3#[cfg(any(test, feature = "test-support"))]
4pub mod test;
5pub mod transport;
6pub mod types;
7
8use std::fmt::Display;
9use std::path::Path;
10use std::sync::Arc;
11
12use anyhow::Result;
13use client::Client;
14use collections::HashMap;
15use gpui::AsyncApp;
16use parking_lot::RwLock;
17use schemars::JsonSchema;
18use serde::{Deserialize, Serialize};
19use util::redact::should_redact;
20
21#[derive(Debug, Clone, PartialEq, Eq, Hash)]
22pub struct ContextServerId(pub Arc<str>);
23
24impl Display for ContextServerId {
25 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26 write!(f, "{}", self.0)
27 }
28}
29
30#[derive(Deserialize, Serialize, Clone, PartialEq, Eq, JsonSchema)]
31pub struct ContextServerCommand {
32 pub path: String,
33 pub args: Vec<String>,
34 pub env: Option<HashMap<String, String>>,
35}
36
37impl std::fmt::Debug for ContextServerCommand {
38 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
39 let filtered_env = self.env.as_ref().map(|env| {
40 env.iter()
41 .map(|(k, v)| (k, if should_redact(k) { "[REDACTED]" } else { v }))
42 .collect::<Vec<_>>()
43 });
44
45 f.debug_struct("ContextServerCommand")
46 .field("path", &self.path)
47 .field("args", &self.args)
48 .field("env", &filtered_env)
49 .finish()
50 }
51}
52
53enum ContextServerTransport {
54 Stdio(ContextServerCommand),
55 Custom(Arc<dyn crate::transport::Transport>),
56}
57
58pub struct ContextServer {
59 id: ContextServerId,
60 client: RwLock<Option<Arc<crate::protocol::InitializedContextServerProtocol>>>,
61 configuration: ContextServerTransport,
62}
63
64impl ContextServer {
65 pub fn stdio(id: ContextServerId, command: ContextServerCommand) -> Self {
66 Self {
67 id,
68 client: RwLock::new(None),
69 configuration: ContextServerTransport::Stdio(command),
70 }
71 }
72
73 pub fn new(id: ContextServerId, transport: Arc<dyn crate::transport::Transport>) -> Self {
74 Self {
75 id,
76 client: RwLock::new(None),
77 configuration: ContextServerTransport::Custom(transport),
78 }
79 }
80
81 pub fn id(&self) -> ContextServerId {
82 self.id.clone()
83 }
84
85 pub fn client(&self) -> Option<Arc<crate::protocol::InitializedContextServerProtocol>> {
86 self.client.read().clone()
87 }
88
89 pub async fn start(self: Arc<Self>, cx: &AsyncApp) -> Result<()> {
90 let client = match &self.configuration {
91 ContextServerTransport::Stdio(command) => Client::stdio(
92 client::ContextServerId(self.id.0.clone()),
93 client::ModelContextServerBinary {
94 executable: Path::new(&command.path).to_path_buf(),
95 args: command.args.clone(),
96 env: command.env.clone(),
97 },
98 cx.clone(),
99 )?,
100 ContextServerTransport::Custom(transport) => Client::new(
101 client::ContextServerId(self.id.0.clone()),
102 self.id().0,
103 transport.clone(),
104 cx.clone(),
105 )?,
106 };
107 self.initialize(client).await
108 }
109
110 async fn initialize(&self, client: Client) -> Result<()> {
111 log::info!("starting context server {}", self.id);
112 let protocol = crate::protocol::ModelContextProtocol::new(client);
113 let client_info = types::Implementation {
114 name: "Zed".to_string(),
115 version: env!("CARGO_PKG_VERSION").to_string(),
116 };
117 let initialized_protocol = protocol.initialize(client_info).await?;
118
119 log::debug!(
120 "context server {} initialized: {:?}",
121 self.id,
122 initialized_protocol.initialize,
123 );
124
125 *self.client.write() = Some(Arc::new(initialized_protocol));
126 Ok(())
127 }
128
129 pub fn stop(&self) -> Result<()> {
130 let mut client = self.client.write();
131 if let Some(protocol) = client.take() {
132 drop(protocol);
133 }
134 Ok(())
135 }
136}