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 smol::channel;
10use std::{
11 fmt,
12 path::PathBuf,
13 pin::pin,
14 sync::{
15 Arc,
16 atomic::{AtomicI32, Ordering::SeqCst},
17 },
18 time::{Duration, Instant},
19};
20use util::{ResultExt, TryFutureExt};
21
22use crate::{
23 transport::{StdioTransport, Transport},
24 types::{CancelledParams, ClientNotification, Notification as _, notifications::Cancelled},
25};
26
27const JSON_RPC_VERSION: &str = "2.0";
28const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
29
30// Standard JSON-RPC error codes
31pub const PARSE_ERROR: i32 = -32700;
32pub const INVALID_REQUEST: i32 = -32600;
33pub const METHOD_NOT_FOUND: i32 = -32601;
34pub const INVALID_PARAMS: i32 = -32602;
35pub const INTERNAL_ERROR: i32 = -32603;
36
37type ResponseHandler = Box<dyn Send + FnOnce(Result<String, Error>)>;
38type NotificationHandler = Box<dyn Send + FnMut(Value, AsyncApp)>;
39type RequestHandler = Box<dyn Send + FnMut(RequestId, &RawValue, AsyncApp)>;
40
41#[derive(Debug, Clone, Eq, PartialEq, Hash, Serialize, Deserialize)]
42#[serde(untagged)]
43pub enum RequestId {
44 Int(i32),
45 Str(String),
46}
47
48pub(crate) struct Client {
49 server_id: ContextServerId,
50 next_id: AtomicI32,
51 outbound_tx: channel::Sender<String>,
52 name: Arc<str>,
53 notification_handlers: Arc<Mutex<HashMap<&'static str, NotificationHandler>>>,
54 response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
55 #[allow(clippy::type_complexity)]
56 #[allow(dead_code)]
57 io_tasks: Mutex<Option<(Task<Option<()>>, Task<Option<()>>)>>,
58 #[allow(dead_code)]
59 output_done_rx: Mutex<Option<barrier::Receiver>>,
60 executor: BackgroundExecutor,
61 #[allow(dead_code)]
62 transport: Arc<dyn Transport>,
63 request_timeout: Option<Duration>,
64}
65
66#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
67#[repr(transparent)]
68pub(crate) struct ContextServerId(pub Arc<str>);
69
70fn is_null_value<T: Serialize>(value: &T) -> bool {
71 matches!(serde_json::to_value(value), Ok(Value::Null))
72}
73
74#[derive(Serialize, Deserialize)]
75pub struct Request<'a, T> {
76 pub jsonrpc: &'static str,
77 pub id: RequestId,
78 pub method: &'a str,
79 #[serde(skip_serializing_if = "is_null_value")]
80 pub params: T,
81}
82
83#[derive(Serialize, Deserialize)]
84pub struct AnyRequest<'a> {
85 pub jsonrpc: &'a str,
86 pub id: RequestId,
87 pub method: &'a str,
88 #[serde(skip_serializing_if = "is_null_value")]
89 pub params: Option<&'a RawValue>,
90}
91
92#[derive(Serialize, Deserialize)]
93struct AnyResponse<'a> {
94 jsonrpc: &'a str,
95 id: RequestId,
96 #[serde(default)]
97 error: Option<Error>,
98 #[serde(borrow)]
99 result: Option<&'a RawValue>,
100}
101
102#[derive(Serialize, Deserialize)]
103#[allow(dead_code)]
104pub(crate) struct Response<T> {
105 pub jsonrpc: &'static str,
106 pub id: RequestId,
107 #[serde(flatten)]
108 pub value: CspResult<T>,
109}
110
111#[derive(Serialize, Deserialize)]
112#[serde(rename_all = "snake_case")]
113pub(crate) enum CspResult<T> {
114 #[serde(rename = "result")]
115 Ok(Option<T>),
116 #[allow(dead_code)]
117 Error(Option<Error>),
118}
119
120#[derive(Serialize, Deserialize)]
121struct Notification<'a, T> {
122 jsonrpc: &'static str,
123 #[serde(borrow)]
124 method: &'a str,
125 params: T,
126}
127
128#[derive(Debug, Clone, Deserialize)]
129struct AnyNotification<'a> {
130 #[expect(
131 unused,
132 reason = "Part of the JSON-RPC protocol - we expect the field to be present in a valid JSON-RPC notification"
133 )]
134 jsonrpc: &'a str,
135 method: String,
136 #[serde(default)]
137 params: Option<Value>,
138}
139
140#[derive(Debug, Serialize, Deserialize)]
141pub(crate) struct Error {
142 pub message: String,
143 pub code: i32,
144}
145
146#[derive(Debug, Clone, Deserialize)]
147pub struct ModelContextServerBinary {
148 pub executable: PathBuf,
149 pub args: Vec<String>,
150 pub env: Option<HashMap<String, String>>,
151 pub timeout: Option<u64>,
152}
153
154impl Client {
155 /// Creates a new Client instance for a context server.
156 ///
157 /// This function initializes a new Client by spawning a child process for the context server,
158 /// setting up communication channels, and initializing handlers for input/output operations.
159 /// It takes a server ID, binary information, and an async app context as input.
160 pub fn stdio(
161 server_id: ContextServerId,
162 binary: ModelContextServerBinary,
163 working_directory: &Option<PathBuf>,
164 cx: AsyncApp,
165 ) -> Result<Self> {
166 log::debug!(
167 "starting context server (executable={:?}, args={:?})",
168 binary.executable,
169 &binary.args
170 );
171
172 let server_name = binary
173 .executable
174 .file_name()
175 .map(|name| name.to_string_lossy().to_string())
176 .unwrap_or_else(String::new);
177
178 let timeout = binary.timeout.map(Duration::from_millis);
179 let transport = Arc::new(StdioTransport::new(binary, working_directory, &cx)?);
180 Self::new(server_id, server_name.into(), transport, timeout, cx)
181 }
182
183 /// Creates a new Client instance for a context server.
184 pub fn new(
185 server_id: ContextServerId,
186 server_name: Arc<str>,
187 transport: Arc<dyn Transport>,
188 request_timeout: Option<Duration>,
189 cx: AsyncApp,
190 ) -> Result<Self> {
191 let (outbound_tx, outbound_rx) = channel::unbounded::<String>();
192 let (output_done_tx, output_done_rx) = barrier::channel();
193
194 let notification_handlers =
195 Arc::new(Mutex::new(HashMap::<_, NotificationHandler>::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 notification_handlers = notification_handlers.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 notification_handlers,
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 notification_handlers,
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 notification_handlers: Arc<Mutex<HashMap<&'static str, NotificationHandler>>>,
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 let mut notification_handlers = notification_handlers.lock();
286 if let Some(handler) = notification_handlers.get_mut(notification.method.as_str()) {
287 handler(notification.params.unwrap_or(Value::Null), cx.clone());
288 }
289 } else {
290 log::error!("Unhandled JSON from context_server: {}", message);
291 }
292 }
293
294 smol::future::yield_now().await;
295
296 Ok(())
297 }
298
299 /// Handles the stderr output from the context server.
300 /// Continuously reads and logs any error messages from the server.
301 async fn handle_err(transport: Arc<dyn Transport>) -> anyhow::Result<()> {
302 while let Some(err) = transport.receive_err().next().await {
303 log::debug!("context server stderr: {}", err.trim());
304 }
305
306 Ok(())
307 }
308
309 /// Handles the output to the context server's stdin.
310 /// This function continuously receives messages from the outbound channel,
311 /// writes them to the server's stdin, and manages the lifecycle of response handlers.
312 async fn handle_output(
313 transport: Arc<dyn Transport>,
314 outbound_rx: channel::Receiver<String>,
315 output_done_tx: barrier::Sender,
316 response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
317 ) -> anyhow::Result<()> {
318 let _clear_response_handlers = util::defer({
319 let response_handlers = response_handlers.clone();
320 move || {
321 response_handlers.lock().take();
322 }
323 });
324 while let Ok(message) = outbound_rx.recv().await {
325 log::trace!("outgoing message: {}", message);
326 transport.send(message).await?;
327 }
328 drop(output_done_tx);
329 Ok(())
330 }
331
332 /// Sends a JSON-RPC request to the context server and waits for a response.
333 /// This function handles serialization, deserialization, timeout, and error handling.
334 pub async fn request<T: DeserializeOwned>(
335 &self,
336 method: &str,
337 params: impl Serialize,
338 ) -> Result<T> {
339 self.request_with(
340 method,
341 params,
342 None,
343 self.request_timeout.or(Some(DEFAULT_REQUEST_TIMEOUT)),
344 )
345 .await
346 }
347
348 pub async fn request_with<T: DeserializeOwned>(
349 &self,
350 method: &str,
351 params: impl Serialize,
352 cancel_rx: Option<oneshot::Receiver<()>>,
353 timeout: Option<Duration>,
354 ) -> Result<T> {
355 let id = self.next_id.fetch_add(1, SeqCst);
356 let request = serde_json::to_string(&Request {
357 jsonrpc: JSON_RPC_VERSION,
358 id: RequestId::Int(id),
359 method,
360 params,
361 })
362 .unwrap();
363
364 let (tx, rx) = oneshot::channel();
365 let handle_response = self
366 .response_handlers
367 .lock()
368 .as_mut()
369 .context("server shut down")
370 .map(|handlers| {
371 handlers.insert(
372 RequestId::Int(id),
373 Box::new(move |result| {
374 let _ = tx.send(result);
375 }),
376 );
377 });
378
379 let send = self
380 .outbound_tx
381 .try_send(request)
382 .context("failed to write to context server's stdin");
383
384 let executor = self.executor.clone();
385 let started = Instant::now();
386 handle_response?;
387 send?;
388
389 let mut timeout_fut = pin!(
390 match timeout {
391 Some(timeout) => future::Either::Left(executor.timer(timeout)),
392 None => future::Either::Right(future::pending()),
393 }
394 .fuse()
395 );
396 let mut cancel_fut = pin!(
397 match cancel_rx {
398 Some(rx) => future::Either::Left(async {
399 rx.await.log_err();
400 }),
401 None => future::Either::Right(future::pending()),
402 }
403 .fuse()
404 );
405
406 select! {
407 response = rx.fuse() => {
408 let elapsed = started.elapsed();
409 log::trace!("took {elapsed:?} to receive response to {method:?} id {id}");
410 match response? {
411 Ok(response) => {
412 let parsed: AnyResponse = serde_json::from_str(&response)?;
413 if let Some(error) = parsed.error {
414 Err(anyhow!(error.message))
415 } else if let Some(result) = parsed.result {
416 Ok(serde_json::from_str(result.get())?)
417 } else {
418 anyhow::bail!("Invalid response: no result or error");
419 }
420 }
421 Err(_) => anyhow::bail!("cancelled")
422 }
423 }
424 _ = cancel_fut => {
425 self.notify(
426 Cancelled::METHOD,
427 ClientNotification::Cancelled(CancelledParams {
428 request_id: RequestId::Int(id),
429 reason: None
430 })
431 ).log_err();
432 anyhow::bail!(RequestCanceled)
433 }
434 _ = timeout_fut => {
435 log::error!("cancelled csp request task for {method:?} id {id} which took over {:?}", timeout.unwrap());
436 anyhow::bail!("Context server request timeout");
437 }
438 }
439 }
440
441 /// Sends a notification to the context server without expecting a response.
442 /// This function serializes the notification and sends it through the outbound channel.
443 pub fn notify(&self, method: &str, params: impl Serialize) -> Result<()> {
444 let notification = serde_json::to_string(&Notification {
445 jsonrpc: JSON_RPC_VERSION,
446 method,
447 params,
448 })
449 .unwrap();
450 self.outbound_tx.try_send(notification)?;
451 Ok(())
452 }
453
454 pub fn on_notification(
455 &self,
456 method: &'static str,
457 f: Box<dyn 'static + Send + FnMut(Value, AsyncApp)>,
458 ) {
459 self.notification_handlers.lock().insert(method, f);
460 }
461}
462
463#[derive(Debug)]
464pub struct RequestCanceled;
465
466impl std::error::Error for RequestCanceled {}
467
468impl std::fmt::Display for RequestCanceled {
469 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
470 f.write_str("Context server request was canceled")
471 }
472}
473
474impl fmt::Display for ContextServerId {
475 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
476 self.0.fmt(f)
477 }
478}
479
480impl fmt::Debug for Client {
481 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
482 f.debug_struct("Context Server Client")
483 .field("id", &self.server_id.0)
484 .field("name", &self.name)
485 .finish_non_exhaustive()
486 }
487}