1use anyhow::{Context as _, Result, anyhow};
2use collections::HashMap;
3use futures::{FutureExt, StreamExt, channel::oneshot, future, select};
4use gpui::{AppContext as _, AsyncApp, BackgroundExecutor, Task};
5use parking_lot::Mutex;
6use postage::barrier;
7use serde::{Deserialize, Serialize, de::DeserializeOwned};
8use serde_json::{Value, value::RawValue};
9use slotmap::SlotMap;
10use smol::channel;
11use std::{
12 fmt,
13 path::PathBuf,
14 pin::pin,
15 sync::{
16 Arc,
17 atomic::{AtomicI32, Ordering::SeqCst},
18 },
19 time::{Duration, Instant},
20};
21use util::{ResultExt, TryFutureExt};
22
23use crate::{
24 transport::{StdioTransport, Transport},
25 types::{CancelledParams, ClientNotification, Notification as _, notifications::Cancelled},
26};
27
28const JSON_RPC_VERSION: &str = "2.0";
29const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
30
31// Standard JSON-RPC error codes
32pub const PARSE_ERROR: i32 = -32700;
33pub const INVALID_REQUEST: i32 = -32600;
34pub const METHOD_NOT_FOUND: i32 = -32601;
35pub const INVALID_PARAMS: i32 = -32602;
36pub const INTERNAL_ERROR: i32 = -32603;
37
38type ResponseHandler = Box<dyn Send + FnOnce(Result<String, Error>)>;
39type NotificationHandler = Box<dyn Send + FnMut(Value, AsyncApp)>;
40type RequestHandler = Box<dyn Send + FnMut(RequestId, &RawValue, AsyncApp)>;
41
42#[derive(Debug, Clone, Eq, PartialEq, Hash, Serialize, Deserialize)]
43#[serde(untagged)]
44pub enum RequestId {
45 Int(i32),
46 Str(String),
47}
48
49pub(crate) struct Client {
50 server_id: ContextServerId,
51 next_id: AtomicI32,
52 outbound_tx: channel::Sender<String>,
53 name: Arc<str>,
54 subscription_set: Arc<Mutex<NotificationSubscriptionSet>>,
55 response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
56 #[allow(clippy::type_complexity)]
57 #[allow(dead_code)]
58 io_tasks: Mutex<Option<(Task<Option<()>>, Task<Option<()>>)>>,
59 #[allow(dead_code)]
60 output_done_rx: Mutex<Option<barrier::Receiver>>,
61 executor: BackgroundExecutor,
62 #[allow(dead_code)]
63 transport: Arc<dyn Transport>,
64 request_timeout: Option<Duration>,
65}
66
67#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
68#[repr(transparent)]
69pub(crate) struct ContextServerId(pub Arc<str>);
70
71fn is_null_value<T: Serialize>(value: &T) -> bool {
72 matches!(serde_json::to_value(value), Ok(Value::Null))
73}
74
75#[derive(Serialize, Deserialize)]
76pub struct Request<'a, T> {
77 pub jsonrpc: &'static str,
78 pub id: RequestId,
79 pub method: &'a str,
80 #[serde(skip_serializing_if = "is_null_value")]
81 pub params: T,
82}
83
84#[derive(Serialize, Deserialize)]
85pub struct AnyRequest<'a> {
86 pub jsonrpc: &'a str,
87 pub id: RequestId,
88 pub method: &'a str,
89 #[serde(skip_serializing_if = "is_null_value")]
90 pub params: Option<&'a RawValue>,
91}
92
93#[derive(Serialize, Deserialize)]
94struct AnyResponse<'a> {
95 jsonrpc: &'a str,
96 id: RequestId,
97 #[serde(default)]
98 error: Option<Error>,
99 #[serde(borrow)]
100 result: Option<&'a RawValue>,
101}
102
103#[derive(Serialize, Deserialize)]
104#[allow(dead_code)]
105pub(crate) struct Response<T> {
106 pub jsonrpc: &'static str,
107 pub id: RequestId,
108 #[serde(flatten)]
109 pub value: CspResult<T>,
110}
111
112#[derive(Serialize, Deserialize)]
113#[serde(rename_all = "snake_case")]
114pub(crate) enum CspResult<T> {
115 #[serde(rename = "result")]
116 Ok(Option<T>),
117 #[allow(dead_code)]
118 Error(Option<Error>),
119}
120
121#[derive(Serialize, Deserialize)]
122struct Notification<'a, T> {
123 jsonrpc: &'static str,
124 #[serde(borrow)]
125 method: &'a str,
126 params: T,
127}
128
129#[derive(Debug, Clone, Deserialize)]
130struct AnyNotification<'a> {
131 #[expect(
132 unused,
133 reason = "Part of the JSON-RPC protocol - we expect the field to be present in a valid JSON-RPC notification"
134 )]
135 jsonrpc: &'a str,
136 method: String,
137 #[serde(default)]
138 params: Option<Value>,
139}
140
141#[derive(Debug, Serialize, Deserialize)]
142pub(crate) struct Error {
143 pub message: String,
144 pub code: i32,
145}
146
147#[derive(Debug, Clone, Deserialize)]
148pub struct ModelContextServerBinary {
149 pub executable: PathBuf,
150 pub args: Vec<String>,
151 pub env: Option<HashMap<String, String>>,
152 pub timeout: Option<u64>,
153}
154
155impl Client {
156 /// Creates a new Client instance for a context server.
157 ///
158 /// This function initializes a new Client by spawning a child process for the context server,
159 /// setting up communication channels, and initializing handlers for input/output operations.
160 /// It takes a server ID, binary information, and an async app context as input.
161 pub fn stdio(
162 server_id: ContextServerId,
163 binary: ModelContextServerBinary,
164 working_directory: &Option<PathBuf>,
165 cx: AsyncApp,
166 ) -> Result<Self> {
167 log::debug!(
168 "starting context server (executable={:?}, args={:?})",
169 binary.executable,
170 &binary.args
171 );
172
173 let server_name = binary
174 .executable
175 .file_name()
176 .map(|name| name.to_string_lossy().into_owned())
177 .unwrap_or_else(String::new);
178
179 let timeout = binary.timeout.map(Duration::from_millis);
180 let transport = Arc::new(StdioTransport::new(binary, working_directory, &cx)?);
181 Self::new(server_id, server_name.into(), transport, timeout, cx)
182 }
183
184 /// Creates a new Client instance for a context server.
185 pub fn new(
186 server_id: ContextServerId,
187 server_name: Arc<str>,
188 transport: Arc<dyn Transport>,
189 request_timeout: Option<Duration>,
190 cx: AsyncApp,
191 ) -> Result<Self> {
192 let (outbound_tx, outbound_rx) = channel::unbounded::<String>();
193 let (output_done_tx, output_done_rx) = barrier::channel();
194
195 let subscription_set = Arc::new(Mutex::new(NotificationSubscriptionSet::default()));
196 let response_handlers =
197 Arc::new(Mutex::new(Some(HashMap::<_, ResponseHandler>::default())));
198 let request_handlers = Arc::new(Mutex::new(HashMap::<_, RequestHandler>::default()));
199
200 let receive_input_task = cx.spawn({
201 let subscription_set = subscription_set.clone();
202 let response_handlers = response_handlers.clone();
203 let request_handlers = request_handlers.clone();
204 let transport = transport.clone();
205 async move |cx| {
206 Self::handle_input(
207 transport,
208 subscription_set,
209 request_handlers,
210 response_handlers,
211 cx,
212 )
213 .log_err()
214 .await
215 }
216 });
217 let receive_err_task = cx.spawn({
218 let transport = transport.clone();
219 async move |_| Self::handle_err(transport).log_err().await
220 });
221 let input_task = cx.spawn(async move |_| {
222 let (input, err) = futures::join!(receive_input_task, receive_err_task);
223 input.or(err)
224 });
225
226 let output_task = cx.background_spawn({
227 let transport = transport.clone();
228 Self::handle_output(
229 transport,
230 outbound_rx,
231 output_done_tx,
232 response_handlers.clone(),
233 )
234 .log_err()
235 });
236
237 Ok(Self {
238 server_id,
239 subscription_set,
240 response_handlers,
241 name: server_name,
242 next_id: Default::default(),
243 outbound_tx,
244 executor: cx.background_executor().clone(),
245 io_tasks: Mutex::new(Some((input_task, output_task))),
246 output_done_rx: Mutex::new(Some(output_done_rx)),
247 transport,
248 request_timeout,
249 })
250 }
251
252 /// Handles input from the server's stdout.
253 ///
254 /// This function continuously reads lines from the provided stdout stream,
255 /// parses them as JSON-RPC responses or notifications, and dispatches them
256 /// to the appropriate handlers. It processes both responses (which are matched
257 /// to pending requests) and notifications (which trigger registered handlers).
258 async fn handle_input(
259 transport: Arc<dyn Transport>,
260 subscription_set: Arc<Mutex<NotificationSubscriptionSet>>,
261 request_handlers: Arc<Mutex<HashMap<&'static str, RequestHandler>>>,
262 response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
263 cx: &mut AsyncApp,
264 ) -> anyhow::Result<()> {
265 let mut receiver = transport.receive();
266
267 while let Some(message) = receiver.next().await {
268 log::trace!("recv: {}", &message);
269 if let Ok(request) = serde_json::from_str::<AnyRequest>(&message) {
270 let mut request_handlers = request_handlers.lock();
271 if let Some(handler) = request_handlers.get_mut(request.method) {
272 handler(
273 request.id,
274 request.params.unwrap_or(RawValue::NULL),
275 cx.clone(),
276 );
277 }
278 } else if let Ok(response) = serde_json::from_str::<AnyResponse>(&message) {
279 if let Some(handlers) = response_handlers.lock().as_mut()
280 && let Some(handler) = handlers.remove(&response.id)
281 {
282 handler(Ok(message.to_string()));
283 }
284 } else if let Ok(notification) = serde_json::from_str::<AnyNotification>(&message) {
285 subscription_set.lock().notify(
286 ¬ification.method,
287 notification.params.unwrap_or(Value::Null),
288 cx,
289 )
290 } else {
291 log::error!("Unhandled JSON from context_server: {}", message);
292 }
293 }
294
295 smol::future::yield_now().await;
296
297 Ok(())
298 }
299
300 /// Handles the stderr output from the context server.
301 /// Continuously reads and logs any error messages from the server.
302 async fn handle_err(transport: Arc<dyn Transport>) -> anyhow::Result<()> {
303 while let Some(err) = transport.receive_err().next().await {
304 log::debug!("context server stderr: {}", err.trim());
305 }
306
307 Ok(())
308 }
309
310 /// Handles the output to the context server's stdin.
311 /// This function continuously receives messages from the outbound channel,
312 /// writes them to the server's stdin, and manages the lifecycle of response handlers.
313 async fn handle_output(
314 transport: Arc<dyn Transport>,
315 outbound_rx: channel::Receiver<String>,
316 output_done_tx: barrier::Sender,
317 response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
318 ) -> anyhow::Result<()> {
319 let _clear_response_handlers = util::defer({
320 let response_handlers = response_handlers.clone();
321 move || {
322 response_handlers.lock().take();
323 }
324 });
325 while let Ok(message) = outbound_rx.recv().await {
326 log::trace!("outgoing message: {}", message);
327 transport.send(message).await?;
328 }
329 drop(output_done_tx);
330 Ok(())
331 }
332
333 /// Sends a JSON-RPC request to the context server and waits for a response.
334 /// This function handles serialization, deserialization, timeout, and error handling.
335 pub async fn request<T: DeserializeOwned>(
336 &self,
337 method: &str,
338 params: impl Serialize,
339 ) -> Result<T> {
340 self.request_with(
341 method,
342 params,
343 None,
344 self.request_timeout.or(Some(DEFAULT_REQUEST_TIMEOUT)),
345 )
346 .await
347 }
348
349 pub async fn request_with<T: DeserializeOwned>(
350 &self,
351 method: &str,
352 params: impl Serialize,
353 cancel_rx: Option<oneshot::Receiver<()>>,
354 timeout: Option<Duration>,
355 ) -> Result<T> {
356 let id = self.next_id.fetch_add(1, SeqCst);
357 let request = serde_json::to_string(&Request {
358 jsonrpc: JSON_RPC_VERSION,
359 id: RequestId::Int(id),
360 method,
361 params,
362 })
363 .unwrap();
364
365 let (tx, rx) = oneshot::channel();
366 let handle_response = self
367 .response_handlers
368 .lock()
369 .as_mut()
370 .context("server shut down")
371 .map(|handlers| {
372 handlers.insert(
373 RequestId::Int(id),
374 Box::new(move |result| {
375 let _ = tx.send(result);
376 }),
377 );
378 });
379
380 let send = self
381 .outbound_tx
382 .try_send(request)
383 .context("failed to write to context server's stdin");
384
385 let executor = self.executor.clone();
386 let started = Instant::now();
387 handle_response?;
388 send?;
389
390 let mut timeout_fut = pin!(
391 match timeout {
392 Some(timeout) => future::Either::Left(executor.timer(timeout)),
393 None => future::Either::Right(future::pending()),
394 }
395 .fuse()
396 );
397 let mut cancel_fut = pin!(
398 match cancel_rx {
399 Some(rx) => future::Either::Left(async {
400 rx.await.log_err();
401 }),
402 None => future::Either::Right(future::pending()),
403 }
404 .fuse()
405 );
406
407 select! {
408 response = rx.fuse() => {
409 let elapsed = started.elapsed();
410 log::trace!("took {elapsed:?} to receive response to {method:?} id {id}");
411 match response? {
412 Ok(response) => {
413 let parsed: AnyResponse = serde_json::from_str(&response)?;
414 if let Some(error) = parsed.error {
415 Err(anyhow!(error.message))
416 } else if let Some(result) = parsed.result {
417 Ok(serde_json::from_str(result.get())?)
418 } else {
419 anyhow::bail!("Invalid response: no result or error");
420 }
421 }
422 Err(_) => anyhow::bail!("cancelled")
423 }
424 }
425 _ = cancel_fut => {
426 self.notify(
427 Cancelled::METHOD,
428 ClientNotification::Cancelled(CancelledParams {
429 request_id: RequestId::Int(id),
430 reason: None
431 })
432 ).log_err();
433 anyhow::bail!(RequestCanceled)
434 }
435 _ = timeout_fut => {
436 log::error!("cancelled csp request task for {method:?} id {id} which took over {:?}", timeout.unwrap());
437 anyhow::bail!("Context server request timeout");
438 }
439 }
440 }
441
442 /// Sends a notification to the context server without expecting a response.
443 /// This function serializes the notification and sends it through the outbound channel.
444 pub fn notify(&self, method: &str, params: impl Serialize) -> Result<()> {
445 let notification = serde_json::to_string(&Notification {
446 jsonrpc: JSON_RPC_VERSION,
447 method,
448 params,
449 })
450 .unwrap();
451 self.outbound_tx.try_send(notification)?;
452 Ok(())
453 }
454
455 #[must_use]
456 pub fn on_notification(
457 &self,
458 method: &'static str,
459 f: Box<dyn 'static + Send + FnMut(Value, AsyncApp)>,
460 ) -> NotificationSubscription {
461 let mut notification_subscriptions = self.subscription_set.lock();
462
463 NotificationSubscription {
464 id: notification_subscriptions.add_handler(method, f),
465 set: self.subscription_set.clone(),
466 }
467 }
468}
469
470#[derive(Debug)]
471pub struct RequestCanceled;
472
473impl std::error::Error for RequestCanceled {}
474
475impl std::fmt::Display for RequestCanceled {
476 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
477 f.write_str("Context server request was canceled")
478 }
479}
480
481impl fmt::Display for ContextServerId {
482 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
483 self.0.fmt(f)
484 }
485}
486
487impl fmt::Debug for Client {
488 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
489 f.debug_struct("Context Server Client")
490 .field("id", &self.server_id.0)
491 .field("name", &self.name)
492 .finish_non_exhaustive()
493 }
494}
495
496slotmap::new_key_type! {
497 struct NotificationSubscriptionId;
498}
499
500#[derive(Default)]
501pub struct NotificationSubscriptionSet {
502 // we have very few subscriptions at the moment
503 methods: Vec<(&'static str, Vec<NotificationSubscriptionId>)>,
504 handlers: SlotMap<NotificationSubscriptionId, NotificationHandler>,
505}
506
507impl NotificationSubscriptionSet {
508 #[must_use]
509 fn add_handler(
510 &mut self,
511 method: &'static str,
512 handler: NotificationHandler,
513 ) -> NotificationSubscriptionId {
514 let id = self.handlers.insert(handler);
515 if let Some((_, handler_ids)) = self
516 .methods
517 .iter_mut()
518 .find(|(probe_method, _)| method == *probe_method)
519 {
520 debug_assert!(
521 handler_ids.len() < 20,
522 "Too many MCP handlers for {}. Consider using a different data structure.",
523 method
524 );
525
526 handler_ids.push(id);
527 } else {
528 self.methods.push((method, vec![id]));
529 };
530 id
531 }
532
533 fn notify(&mut self, method: &str, payload: Value, cx: &mut AsyncApp) {
534 let Some((_, handler_ids)) = self
535 .methods
536 .iter_mut()
537 .find(|(probe_method, _)| method == *probe_method)
538 else {
539 return;
540 };
541
542 for handler_id in handler_ids {
543 if let Some(handler) = self.handlers.get_mut(*handler_id) {
544 handler(payload.clone(), cx.clone());
545 }
546 }
547 }
548}
549
550pub struct NotificationSubscription {
551 id: NotificationSubscriptionId,
552 set: Arc<Mutex<NotificationSubscriptionSet>>,
553}
554
555impl Drop for NotificationSubscription {
556 fn drop(&mut self) {
557 let mut set = self.set.lock();
558 set.handlers.remove(self.id);
559 set.methods.retain_mut(|(_, handler_ids)| {
560 handler_ids.retain(|id| *id != self.id);
561 !handler_ids.is_empty()
562 });
563 }
564}