1use anyhow::{anyhow, Context, Result};
2use collections::HashMap;
3use futures::{channel::oneshot, io::BufWriter, select, AsyncRead, AsyncWrite, FutureExt};
4use gpui::{AsyncAppContext, BackgroundExecutor, Task};
5use parking_lot::Mutex;
6use postage::barrier;
7use serde::{de::DeserializeOwned, Deserialize, Serialize};
8use serde_json::{value::RawValue, Value};
9use smol::{
10 channel,
11 io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
12 process::{self, Child},
13};
14use std::{
15 fmt,
16 path::PathBuf,
17 sync::{
18 atomic::{AtomicI32, Ordering::SeqCst},
19 Arc,
20 },
21 time::{Duration, Instant},
22};
23use util::TryFutureExt;
24
25const JSON_RPC_VERSION: &str = "2.0";
26const REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
27
28type ResponseHandler = Box<dyn Send + FnOnce(Result<String, Error>)>;
29type NotificationHandler = Box<dyn Send + FnMut(RequestId, Value, AsyncAppContext)>;
30
31#[derive(Debug, Clone, Eq, PartialEq, Hash, Serialize, Deserialize)]
32#[serde(untagged)]
33pub enum RequestId {
34 Int(i32),
35 Str(String),
36}
37
38pub struct Client {
39 server_id: ContextServerId,
40 next_id: AtomicI32,
41 outbound_tx: channel::Sender<String>,
42 name: Arc<str>,
43 notification_handlers: Arc<Mutex<HashMap<&'static str, NotificationHandler>>>,
44 response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
45 #[allow(clippy::type_complexity)]
46 #[allow(dead_code)]
47 io_tasks: Mutex<Option<(Task<Option<()>>, Task<Option<()>>)>>,
48 #[allow(dead_code)]
49 output_done_rx: Mutex<Option<barrier::Receiver>>,
50 executor: BackgroundExecutor,
51 server: Arc<Mutex<Option<Child>>>,
52}
53
54#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
55#[repr(transparent)]
56pub struct ContextServerId(pub String);
57
58#[derive(Serialize, Deserialize)]
59struct Request<'a, T> {
60 jsonrpc: &'static str,
61 id: RequestId,
62 method: &'a str,
63 params: T,
64}
65
66#[derive(Serialize, Deserialize)]
67struct AnyResponse<'a> {
68 jsonrpc: &'a str,
69 id: RequestId,
70 #[serde(default)]
71 error: Option<Error>,
72 #[serde(borrow)]
73 result: Option<&'a RawValue>,
74}
75
76#[derive(Deserialize)]
77#[allow(dead_code)]
78struct Response<T> {
79 jsonrpc: &'static str,
80 id: RequestId,
81 #[serde(flatten)]
82 value: CspResult<T>,
83}
84
85#[derive(Deserialize)]
86#[serde(rename_all = "snake_case")]
87enum CspResult<T> {
88 #[serde(rename = "result")]
89 Ok(Option<T>),
90 #[allow(dead_code)]
91 Error(Option<Error>),
92}
93
94#[derive(Serialize, Deserialize)]
95struct Notification<'a, T> {
96 jsonrpc: &'static str,
97 id: RequestId,
98 #[serde(borrow)]
99 method: &'a str,
100 params: T,
101}
102
103#[derive(Debug, Clone, Deserialize)]
104struct AnyNotification<'a> {
105 jsonrpc: &'a str,
106 id: RequestId,
107 method: String,
108 #[serde(default)]
109 params: Option<Value>,
110}
111
112#[derive(Debug, Serialize, Deserialize)]
113struct Error {
114 message: String,
115}
116
117#[derive(Debug, Clone, Deserialize)]
118pub struct ModelContextServerBinary {
119 pub executable: PathBuf,
120 pub args: Vec<String>,
121 pub env: Option<HashMap<String, String>>,
122}
123
124impl Client {
125 /// Creates a new Client instance for a context server.
126 ///
127 /// This function initializes a new Client by spawning a child process for the context server,
128 /// setting up communication channels, and initializing handlers for input/output operations.
129 /// It takes a server ID, binary information, and an async app context as input.
130 pub fn new(
131 server_id: ContextServerId,
132 binary: ModelContextServerBinary,
133 cx: AsyncAppContext,
134 ) -> Result<Self> {
135 log::info!(
136 "starting context server (executable={:?}, args={:?})",
137 binary.executable,
138 &binary.args
139 );
140
141 let mut command = process::Command::new(&binary.executable);
142 command
143 .args(&binary.args)
144 .envs(binary.env.unwrap_or_default())
145 .stdin(std::process::Stdio::piped())
146 .stdout(std::process::Stdio::piped())
147 .stderr(std::process::Stdio::piped())
148 .kill_on_drop(true);
149
150 let mut server = command.spawn().with_context(|| {
151 format!(
152 "failed to spawn command. (path={:?}, args={:?})",
153 binary.executable, &binary.args
154 )
155 })?;
156
157 let stdin = server.stdin.take().unwrap();
158 let stdout = server.stdout.take().unwrap();
159 let stderr = server.stderr.take().unwrap();
160
161 let (outbound_tx, outbound_rx) = channel::unbounded::<String>();
162 let (output_done_tx, output_done_rx) = barrier::channel();
163
164 let notification_handlers =
165 Arc::new(Mutex::new(HashMap::<_, NotificationHandler>::default()));
166 let response_handlers =
167 Arc::new(Mutex::new(Some(HashMap::<_, ResponseHandler>::default())));
168
169 let stdout_input_task = cx.spawn({
170 let notification_handlers = notification_handlers.clone();
171 let response_handlers = response_handlers.clone();
172 move |cx| {
173 Self::handle_input(stdout, notification_handlers, response_handlers, cx).log_err()
174 }
175 });
176 let stderr_input_task = cx.spawn(|_| Self::handle_stderr(stderr).log_err());
177 let input_task = cx.spawn(|_| async move {
178 let (stdout, stderr) = futures::join!(stdout_input_task, stderr_input_task);
179 stdout.or(stderr)
180 });
181 let output_task = cx.background_executor().spawn({
182 Self::handle_output(
183 stdin,
184 outbound_rx,
185 output_done_tx,
186 response_handlers.clone(),
187 )
188 .log_err()
189 });
190
191 let mut context_server = Self {
192 server_id,
193 notification_handlers,
194 response_handlers,
195 name: "".into(),
196 next_id: Default::default(),
197 outbound_tx,
198 executor: cx.background_executor().clone(),
199 io_tasks: Mutex::new(Some((input_task, output_task))),
200 output_done_rx: Mutex::new(Some(output_done_rx)),
201 server: Arc::new(Mutex::new(Some(server))),
202 };
203
204 if let Some(name) = binary.executable.file_name() {
205 context_server.name = name.to_string_lossy().into();
206 }
207
208 Ok(context_server)
209 }
210
211 /// Handles input from the server's stdout.
212 ///
213 /// This function continuously reads lines from the provided stdout stream,
214 /// parses them as JSON-RPC responses or notifications, and dispatches them
215 /// to the appropriate handlers. It processes both responses (which are matched
216 /// to pending requests) and notifications (which trigger registered handlers).
217 async fn handle_input<Stdout>(
218 stdout: Stdout,
219 notification_handlers: Arc<Mutex<HashMap<&'static str, NotificationHandler>>>,
220 response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
221 cx: AsyncAppContext,
222 ) -> anyhow::Result<()>
223 where
224 Stdout: AsyncRead + Unpin + Send + 'static,
225 {
226 let mut stdout = BufReader::new(stdout);
227 let mut buffer = String::new();
228
229 loop {
230 buffer.clear();
231 if stdout.read_line(&mut buffer).await? == 0 {
232 return Ok(());
233 }
234
235 let content = buffer.trim();
236
237 if !content.is_empty() {
238 if let Ok(response) = serde_json::from_str::<AnyResponse>(&content) {
239 if let Some(handlers) = response_handlers.lock().as_mut() {
240 if let Some(handler) = handlers.remove(&response.id) {
241 handler(Ok(content.to_string()));
242 }
243 }
244 } else if let Ok(notification) = serde_json::from_str::<AnyNotification>(&content) {
245 let mut notification_handlers = notification_handlers.lock();
246 if let Some(handler) =
247 notification_handlers.get_mut(notification.method.as_str())
248 {
249 handler(
250 notification.id,
251 notification.params.unwrap_or(Value::Null),
252 cx.clone(),
253 );
254 }
255 }
256 }
257
258 smol::future::yield_now().await;
259 }
260 }
261
262 /// Handles the stderr output from the context server.
263 /// Continuously reads and logs any error messages from the server.
264 async fn handle_stderr<Stderr>(stderr: Stderr) -> anyhow::Result<()>
265 where
266 Stderr: AsyncRead + Unpin + Send + 'static,
267 {
268 let mut stderr = BufReader::new(stderr);
269 let mut buffer = String::new();
270
271 loop {
272 buffer.clear();
273 if stderr.read_line(&mut buffer).await? == 0 {
274 return Ok(());
275 }
276 log::warn!("context server stderr: {}", buffer.trim());
277 smol::future::yield_now().await;
278 }
279 }
280
281 /// Handles the output to the context server's stdin.
282 /// This function continuously receives messages from the outbound channel,
283 /// writes them to the server's stdin, and manages the lifecycle of response handlers.
284 async fn handle_output<Stdin>(
285 stdin: Stdin,
286 outbound_rx: channel::Receiver<String>,
287 output_done_tx: barrier::Sender,
288 response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
289 ) -> anyhow::Result<()>
290 where
291 Stdin: AsyncWrite + Unpin + Send + 'static,
292 {
293 let mut stdin = BufWriter::new(stdin);
294 let _clear_response_handlers = util::defer({
295 let response_handlers = response_handlers.clone();
296 move || {
297 response_handlers.lock().take();
298 }
299 });
300 while let Ok(message) = outbound_rx.recv().await {
301 log::trace!("outgoing message: {}", message);
302
303 stdin.write_all(message.as_bytes()).await?;
304 stdin.write_all(b"\n").await?;
305 stdin.flush().await?;
306 }
307 drop(output_done_tx);
308 Ok(())
309 }
310
311 /// Sends a JSON-RPC request to the context server and waits for a response.
312 /// This function handles serialization, deserialization, timeout, and error handling.
313 pub async fn request<T: DeserializeOwned>(
314 &self,
315 method: &str,
316 params: impl Serialize,
317 ) -> Result<T> {
318 let id = self.next_id.fetch_add(1, SeqCst);
319 let request = serde_json::to_string(&Request {
320 jsonrpc: JSON_RPC_VERSION,
321 id: RequestId::Int(id),
322 method,
323 params,
324 })
325 .unwrap();
326
327 let (tx, rx) = oneshot::channel();
328 let handle_response = self
329 .response_handlers
330 .lock()
331 .as_mut()
332 .ok_or_else(|| anyhow!("server shut down"))
333 .map(|handlers| {
334 handlers.insert(
335 RequestId::Int(id),
336 Box::new(move |result| {
337 let _ = tx.send(result);
338 }),
339 );
340 });
341
342 let send = self
343 .outbound_tx
344 .try_send(request)
345 .context("failed to write to context server's stdin");
346
347 let executor = self.executor.clone();
348 let started = Instant::now();
349 handle_response?;
350 send?;
351
352 let mut timeout = executor.timer(REQUEST_TIMEOUT).fuse();
353 select! {
354 response = rx.fuse() => {
355 let elapsed = started.elapsed();
356 log::trace!("took {elapsed:?} to receive response to {method:?} id {id}");
357 match response? {
358 Ok(response) => {
359 let parsed: AnyResponse = serde_json::from_str(&response)?;
360 if let Some(error) = parsed.error {
361 Err(anyhow!(error.message))
362 } else if let Some(result) = parsed.result {
363 Ok(serde_json::from_str(result.get())?)
364 } else {
365 Err(anyhow!("Invalid response: no result or error"))
366 }
367 }
368 Err(_) => anyhow::bail!("cancelled")
369 }
370 }
371 _ = timeout => {
372 log::error!("cancelled csp request task for {method:?} id {id} which took over {:?}", REQUEST_TIMEOUT);
373 anyhow::bail!("Context server request timeout");
374 }
375 }
376 }
377
378 /// Sends a notification to the context server without expecting a response.
379 /// This function serializes the notification and sends it through the outbound channel.
380 pub fn notify(&self, method: &str, params: impl Serialize) -> Result<()> {
381 let id = self.next_id.fetch_add(1, SeqCst);
382 let notification = serde_json::to_string(&Notification {
383 jsonrpc: JSON_RPC_VERSION,
384 id: RequestId::Int(id),
385 method,
386 params,
387 })
388 .unwrap();
389 self.outbound_tx.try_send(notification)?;
390 Ok(())
391 }
392
393 pub fn on_notification<F>(&self, method: &'static str, mut f: F)
394 where
395 F: 'static + Send + FnMut(Value, AsyncAppContext),
396 {
397 self.notification_handlers
398 .lock()
399 .insert(method, Box::new(move |_, params, cx| f(params, cx)));
400 }
401
402 pub fn name(&self) -> &str {
403 &self.name
404 }
405
406 pub fn server_id(&self) -> ContextServerId {
407 self.server_id.clone()
408 }
409}
410
411impl Drop for Client {
412 fn drop(&mut self) {
413 if let Some(mut server) = self.server.lock().take() {
414 let _ = server.kill();
415 }
416 }
417}
418
419impl fmt::Display for ContextServerId {
420 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
421 self.0.fmt(f)
422 }
423}
424
425impl fmt::Debug for Client {
426 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
427 f.debug_struct("Context Server Client")
428 .field("id", &self.server_id.0)
429 .field("name", &self.name)
430 .finish_non_exhaustive()
431 }
432}