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