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 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}
64
65#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
66#[repr(transparent)]
67pub(crate) struct ContextServerId(pub Arc<str>);
68
69fn is_null_value<T: Serialize>(value: &T) -> bool {
70 matches!(serde_json::to_value(value), Ok(Value::Null))
71}
72
73#[derive(Serialize, Deserialize)]
74pub struct Request<'a, T> {
75 pub jsonrpc: &'static str,
76 pub id: RequestId,
77 pub method: &'a str,
78 #[serde(skip_serializing_if = "is_null_value")]
79 pub params: T,
80}
81
82#[derive(Serialize, Deserialize)]
83pub struct AnyRequest<'a> {
84 pub jsonrpc: &'a str,
85 pub id: RequestId,
86 pub method: &'a str,
87 #[serde(skip_serializing_if = "is_null_value")]
88 pub params: Option<&'a RawValue>,
89}
90
91#[derive(Serialize, Deserialize)]
92struct AnyResponse<'a> {
93 jsonrpc: &'a str,
94 id: RequestId,
95 #[serde(default)]
96 error: Option<Error>,
97 #[serde(borrow)]
98 result: Option<&'a RawValue>,
99}
100
101#[derive(Serialize, Deserialize)]
102#[allow(dead_code)]
103pub(crate) struct Response<T> {
104 pub jsonrpc: &'static str,
105 pub id: RequestId,
106 #[serde(flatten)]
107 pub value: CspResult<T>,
108}
109
110#[derive(Serialize, Deserialize)]
111#[serde(rename_all = "snake_case")]
112pub(crate) enum CspResult<T> {
113 #[serde(rename = "result")]
114 Ok(Option<T>),
115 #[allow(dead_code)]
116 Error(Option<Error>),
117}
118
119#[derive(Serialize, Deserialize)]
120struct Notification<'a, T> {
121 jsonrpc: &'static str,
122 #[serde(borrow)]
123 method: &'a str,
124 params: T,
125}
126
127#[derive(Debug, Clone, Deserialize)]
128struct AnyNotification<'a> {
129 jsonrpc: &'a str,
130 method: String,
131 #[serde(default)]
132 params: Option<Value>,
133}
134
135#[derive(Debug, Serialize, Deserialize)]
136pub(crate) struct Error {
137 pub message: String,
138 pub code: i32,
139}
140
141#[derive(Debug, Clone, Deserialize)]
142pub struct ModelContextServerBinary {
143 pub executable: PathBuf,
144 pub args: Vec<String>,
145 pub env: Option<HashMap<String, String>>,
146}
147
148impl Client {
149 /// Creates a new Client instance for a context server.
150 ///
151 /// This function initializes a new Client by spawning a child process for the context server,
152 /// setting up communication channels, and initializing handlers for input/output operations.
153 /// It takes a server ID, binary information, and an async app context as input.
154 pub fn stdio(
155 server_id: ContextServerId,
156 binary: ModelContextServerBinary,
157 working_directory: &Option<PathBuf>,
158 cx: AsyncApp,
159 ) -> Result<Self> {
160 log::debug!(
161 "starting context server (executable={:?}, args={:?})",
162 binary.executable,
163 &binary.args
164 );
165
166 let server_name = binary
167 .executable
168 .file_name()
169 .map(|name| name.to_string_lossy().to_string())
170 .unwrap_or_else(String::new);
171
172 let transport = Arc::new(StdioTransport::new(binary, working_directory, &cx)?);
173 Self::new(server_id, server_name.into(), transport, cx)
174 }
175
176 /// Creates a new Client instance for a context server.
177 pub fn new(
178 server_id: ContextServerId,
179 server_name: Arc<str>,
180 transport: Arc<dyn Transport>,
181 cx: AsyncApp,
182 ) -> Result<Self> {
183 let (outbound_tx, outbound_rx) = channel::unbounded::<String>();
184 let (output_done_tx, output_done_rx) = barrier::channel();
185
186 let notification_handlers =
187 Arc::new(Mutex::new(HashMap::<_, NotificationHandler>::default()));
188 let response_handlers =
189 Arc::new(Mutex::new(Some(HashMap::<_, ResponseHandler>::default())));
190 let request_handlers = Arc::new(Mutex::new(HashMap::<_, RequestHandler>::default()));
191
192 let receive_input_task = cx.spawn({
193 let notification_handlers = notification_handlers.clone();
194 let response_handlers = response_handlers.clone();
195 let request_handlers = request_handlers.clone();
196 let transport = transport.clone();
197 async move |cx| {
198 Self::handle_input(
199 transport,
200 notification_handlers,
201 request_handlers,
202 response_handlers,
203 cx,
204 )
205 .log_err()
206 .await
207 }
208 });
209 let receive_err_task = cx.spawn({
210 let transport = transport.clone();
211 async move |_| Self::handle_err(transport).log_err().await
212 });
213 let input_task = cx.spawn(async move |_| {
214 let (input, err) = futures::join!(receive_input_task, receive_err_task);
215 input.or(err)
216 });
217
218 let output_task = cx.background_spawn({
219 let transport = transport.clone();
220 Self::handle_output(
221 transport,
222 outbound_rx,
223 output_done_tx,
224 response_handlers.clone(),
225 )
226 .log_err()
227 });
228
229 Ok(Self {
230 server_id,
231 notification_handlers,
232 response_handlers,
233 name: server_name,
234 next_id: Default::default(),
235 outbound_tx,
236 executor: cx.background_executor().clone(),
237 io_tasks: Mutex::new(Some((input_task, output_task))),
238 output_done_rx: Mutex::new(Some(output_done_rx)),
239 transport,
240 })
241 }
242
243 /// Handles input from the server's stdout.
244 ///
245 /// This function continuously reads lines from the provided stdout stream,
246 /// parses them as JSON-RPC responses or notifications, and dispatches them
247 /// to the appropriate handlers. It processes both responses (which are matched
248 /// to pending requests) and notifications (which trigger registered handlers).
249 async fn handle_input(
250 transport: Arc<dyn Transport>,
251 notification_handlers: Arc<Mutex<HashMap<&'static str, NotificationHandler>>>,
252 request_handlers: Arc<Mutex<HashMap<&'static str, RequestHandler>>>,
253 response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
254 cx: &mut AsyncApp,
255 ) -> anyhow::Result<()> {
256 let mut receiver = transport.receive();
257
258 while let Some(message) = receiver.next().await {
259 log::trace!("recv: {}", &message);
260 if let Ok(request) = serde_json::from_str::<AnyRequest>(&message) {
261 let mut request_handlers = request_handlers.lock();
262 if let Some(handler) = request_handlers.get_mut(request.method) {
263 handler(
264 request.id,
265 request.params.unwrap_or(RawValue::NULL),
266 cx.clone(),
267 );
268 }
269 } else if let Ok(response) = serde_json::from_str::<AnyResponse>(&message) {
270 if let Some(handlers) = response_handlers.lock().as_mut()
271 && let Some(handler) = handlers.remove(&response.id)
272 {
273 handler(Ok(message.to_string()));
274 }
275 } else if let Ok(notification) = serde_json::from_str::<AnyNotification>(&message) {
276 let mut notification_handlers = notification_handlers.lock();
277 if let Some(handler) = notification_handlers.get_mut(notification.method.as_str()) {
278 handler(notification.params.unwrap_or(Value::Null), cx.clone());
279 }
280 } else {
281 log::error!("Unhandled JSON from context_server: {}", message);
282 }
283 }
284
285 smol::future::yield_now().await;
286
287 Ok(())
288 }
289
290 /// Handles the stderr output from the context server.
291 /// Continuously reads and logs any error messages from the server.
292 async fn handle_err(transport: Arc<dyn Transport>) -> anyhow::Result<()> {
293 while let Some(err) = transport.receive_err().next().await {
294 log::debug!("context server stderr: {}", err.trim());
295 }
296
297 Ok(())
298 }
299
300 /// Handles the output to the context server's stdin.
301 /// This function continuously receives messages from the outbound channel,
302 /// writes them to the server's stdin, and manages the lifecycle of response handlers.
303 async fn handle_output(
304 transport: Arc<dyn Transport>,
305 outbound_rx: channel::Receiver<String>,
306 output_done_tx: barrier::Sender,
307 response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
308 ) -> anyhow::Result<()> {
309 let _clear_response_handlers = util::defer({
310 let response_handlers = response_handlers.clone();
311 move || {
312 response_handlers.lock().take();
313 }
314 });
315 while let Ok(message) = outbound_rx.recv().await {
316 log::trace!("outgoing message: {}", message);
317 transport.send(message).await?;
318 }
319 drop(output_done_tx);
320 Ok(())
321 }
322
323 /// Sends a JSON-RPC request to the context server and waits for a response.
324 /// This function handles serialization, deserialization, timeout, and error handling.
325 pub async fn request<T: DeserializeOwned>(
326 &self,
327 method: &str,
328 params: impl Serialize,
329 ) -> Result<T> {
330 self.request_with(method, params, None, Some(REQUEST_TIMEOUT))
331 .await
332 }
333
334 pub async fn request_with<T: DeserializeOwned>(
335 &self,
336 method: &str,
337 params: impl Serialize,
338 cancel_rx: Option<oneshot::Receiver<()>>,
339 timeout: Option<Duration>,
340 ) -> Result<T> {
341 let id = self.next_id.fetch_add(1, SeqCst);
342 let request = serde_json::to_string(&Request {
343 jsonrpc: JSON_RPC_VERSION,
344 id: RequestId::Int(id),
345 method,
346 params,
347 })
348 .unwrap();
349
350 let (tx, rx) = oneshot::channel();
351 let handle_response = self
352 .response_handlers
353 .lock()
354 .as_mut()
355 .context("server shut down")
356 .map(|handlers| {
357 handlers.insert(
358 RequestId::Int(id),
359 Box::new(move |result| {
360 let _ = tx.send(result);
361 }),
362 );
363 });
364
365 let send = self
366 .outbound_tx
367 .try_send(request)
368 .context("failed to write to context server's stdin");
369
370 let executor = self.executor.clone();
371 let started = Instant::now();
372 handle_response?;
373 send?;
374
375 let mut timeout_fut = pin!(
376 match timeout {
377 Some(timeout) => future::Either::Left(executor.timer(timeout)),
378 None => future::Either::Right(future::pending()),
379 }
380 .fuse()
381 );
382 let mut cancel_fut = pin!(
383 match cancel_rx {
384 Some(rx) => future::Either::Left(async {
385 rx.await.log_err();
386 }),
387 None => future::Either::Right(future::pending()),
388 }
389 .fuse()
390 );
391
392 select! {
393 response = rx.fuse() => {
394 let elapsed = started.elapsed();
395 log::trace!("took {elapsed:?} to receive response to {method:?} id {id}");
396 match response? {
397 Ok(response) => {
398 let parsed: AnyResponse = serde_json::from_str(&response)?;
399 if let Some(error) = parsed.error {
400 Err(anyhow!(error.message))
401 } else if let Some(result) = parsed.result {
402 Ok(serde_json::from_str(result.get())?)
403 } else {
404 anyhow::bail!("Invalid response: no result or error");
405 }
406 }
407 Err(_) => anyhow::bail!("cancelled")
408 }
409 }
410 _ = cancel_fut => {
411 self.notify(
412 Cancelled::METHOD,
413 ClientNotification::Cancelled(CancelledParams {
414 request_id: RequestId::Int(id),
415 reason: None
416 })
417 ).log_err();
418 anyhow::bail!(RequestCanceled)
419 }
420 _ = timeout_fut => {
421 log::error!("cancelled csp request task for {method:?} id {id} which took over {:?}", timeout.unwrap());
422 anyhow::bail!("Context server request timeout");
423 }
424 }
425 }
426
427 /// Sends a notification to the context server without expecting a response.
428 /// This function serializes the notification and sends it through the outbound channel.
429 pub fn notify(&self, method: &str, params: impl Serialize) -> Result<()> {
430 let notification = serde_json::to_string(&Notification {
431 jsonrpc: JSON_RPC_VERSION,
432 method,
433 params,
434 })
435 .unwrap();
436 self.outbound_tx.try_send(notification)?;
437 Ok(())
438 }
439
440 pub fn on_notification(
441 &self,
442 method: &'static str,
443 f: Box<dyn 'static + Send + FnMut(Value, AsyncApp)>,
444 ) {
445 self.notification_handlers.lock().insert(method, f);
446 }
447}
448
449#[derive(Debug)]
450pub struct RequestCanceled;
451
452impl std::error::Error for RequestCanceled {}
453
454impl std::fmt::Display for RequestCanceled {
455 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
456 f.write_str("Context server request was canceled")
457 }
458}
459
460impl fmt::Display for ContextServerId {
461 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
462 self.0.fmt(f)
463 }
464}
465
466impl fmt::Debug for Client {
467 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
468 f.debug_struct("Context Server Client")
469 .field("id", &self.server_id.0)
470 .field("name", &self.name)
471 .finish_non_exhaustive()
472 }
473}