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 request_handlers: Arc<Mutex<HashMap<&'static str, RequestHandler>>>,
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}
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 if let Ok(Value::Null) = serde_json::to_value(value) {
72 true
73 } else {
74 false
75 }
76}
77
78#[derive(Serialize, Deserialize)]
79pub struct Request<'a, T> {
80 pub jsonrpc: &'static str,
81 pub id: RequestId,
82 pub method: &'a str,
83 #[serde(skip_serializing_if = "is_null_value")]
84 pub params: T,
85}
86
87#[derive(Serialize, Deserialize)]
88pub struct AnyRequest<'a> {
89 pub jsonrpc: &'a str,
90 pub id: RequestId,
91 pub method: &'a str,
92 #[serde(skip_serializing_if = "is_null_value")]
93 pub params: Option<&'a RawValue>,
94}
95
96#[derive(Serialize, Deserialize)]
97struct AnyResponse<'a> {
98 jsonrpc: &'a str,
99 id: RequestId,
100 #[serde(default)]
101 error: Option<Error>,
102 #[serde(borrow)]
103 result: Option<&'a RawValue>,
104}
105
106#[derive(Serialize, Deserialize)]
107#[allow(dead_code)]
108pub(crate) struct Response<T> {
109 pub jsonrpc: &'static str,
110 pub id: RequestId,
111 #[serde(flatten)]
112 pub value: CspResult<T>,
113}
114
115#[derive(Serialize, Deserialize)]
116#[serde(rename_all = "snake_case")]
117pub(crate) enum CspResult<T> {
118 #[serde(rename = "result")]
119 Ok(Option<T>),
120 #[allow(dead_code)]
121 Error(Option<Error>),
122}
123
124#[derive(Serialize, Deserialize)]
125struct Notification<'a, T> {
126 jsonrpc: &'static str,
127 #[serde(borrow)]
128 method: &'a str,
129 params: T,
130}
131
132#[derive(Debug, Clone, Deserialize)]
133struct AnyNotification<'a> {
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}
152
153impl Client {
154 /// Creates a new Client instance for a context server.
155 ///
156 /// This function initializes a new Client by spawning a child process for the context server,
157 /// setting up communication channels, and initializing handlers for input/output operations.
158 /// It takes a server ID, binary information, and an async app context as input.
159 pub fn stdio(
160 server_id: ContextServerId,
161 binary: ModelContextServerBinary,
162 cx: AsyncApp,
163 ) -> Result<Self> {
164 log::info!(
165 "starting context server (executable={:?}, args={:?})",
166 binary.executable,
167 &binary.args
168 );
169
170 let server_name = binary
171 .executable
172 .file_name()
173 .map(|name| name.to_string_lossy().to_string())
174 .unwrap_or_else(String::new);
175
176 let transport = Arc::new(StdioTransport::new(binary, &cx)?);
177 Self::new(server_id, server_name.into(), transport, cx)
178 }
179
180 /// Creates a new Client instance for a context server.
181 pub fn new(
182 server_id: ContextServerId,
183 server_name: Arc<str>,
184 transport: Arc<dyn Transport>,
185 cx: AsyncApp,
186 ) -> Result<Self> {
187 let (outbound_tx, outbound_rx) = channel::unbounded::<String>();
188 let (output_done_tx, output_done_rx) = barrier::channel();
189
190 let notification_handlers =
191 Arc::new(Mutex::new(HashMap::<_, NotificationHandler>::default()));
192 let response_handlers =
193 Arc::new(Mutex::new(Some(HashMap::<_, ResponseHandler>::default())));
194 let request_handlers = Arc::new(Mutex::new(HashMap::<_, RequestHandler>::default()));
195
196 let receive_input_task = cx.spawn({
197 let notification_handlers = notification_handlers.clone();
198 let response_handlers = response_handlers.clone();
199 let request_handlers = request_handlers.clone();
200 let transport = transport.clone();
201 async move |cx| {
202 Self::handle_input(
203 transport,
204 notification_handlers,
205 request_handlers,
206 response_handlers,
207 cx,
208 )
209 .log_err()
210 .await
211 }
212 });
213 let receive_err_task = cx.spawn({
214 let transport = transport.clone();
215 async move |_| Self::handle_err(transport).log_err().await
216 });
217 let input_task = cx.spawn(async move |_| {
218 let (input, err) = futures::join!(receive_input_task, receive_err_task);
219 input.or(err)
220 });
221
222 let output_task = cx.background_spawn({
223 let transport = transport.clone();
224 Self::handle_output(
225 transport,
226 outbound_rx,
227 output_done_tx,
228 response_handlers.clone(),
229 )
230 .log_err()
231 });
232
233 Ok(Self {
234 server_id,
235 notification_handlers,
236 response_handlers,
237 request_handlers,
238 name: server_name,
239 next_id: Default::default(),
240 outbound_tx,
241 executor: cx.background_executor().clone(),
242 io_tasks: Mutex::new(Some((input_task, output_task))),
243 output_done_rx: Mutex::new(Some(output_done_rx)),
244 transport,
245 })
246 }
247
248 /// Handles input from the server's stdout.
249 ///
250 /// This function continuously reads lines from the provided stdout stream,
251 /// parses them as JSON-RPC responses or notifications, and dispatches them
252 /// to the appropriate handlers. It processes both responses (which are matched
253 /// to pending requests) and notifications (which trigger registered handlers).
254 async fn handle_input(
255 transport: Arc<dyn Transport>,
256 notification_handlers: Arc<Mutex<HashMap<&'static str, NotificationHandler>>>,
257 request_handlers: Arc<Mutex<HashMap<&'static str, RequestHandler>>>,
258 response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
259 cx: &mut AsyncApp,
260 ) -> anyhow::Result<()> {
261 let mut receiver = transport.receive();
262
263 while let Some(message) = receiver.next().await {
264 log::trace!("recv: {}", &message);
265 if let Ok(response) = serde_json::from_str::<AnyResponse>(&message) {
266 dbg!("here!");
267 if let Some(handlers) = response_handlers.lock().as_mut() {
268 if let Some(handler) = handlers.remove(&response.id) {
269 handler(Ok(message.to_string()));
270 }
271 }
272 } else if let Some(request) = serde_json::from_str::<AnyRequest>(&message).log_err() {
273 dbg!("here!");
274 let mut request_handlers = request_handlers.lock();
275 if let Some(handler) = request_handlers.get_mut(request.method) {
276 handler(
277 request.id,
278 request.params.unwrap_or(RawValue::NULL),
279 cx.clone(),
280 );
281 }
282 } else if let Ok(notification) = serde_json::from_str::<AnyNotification>(&message) {
283 dbg!("here!");
284 let mut notification_handlers = notification_handlers.lock();
285 if let Some(handler) = notification_handlers.get_mut(notification.method.as_str()) {
286 handler(notification.params.unwrap_or(Value::Null), cx.clone());
287 }
288 } else {
289 dbg!("WTF", &message);
290 }
291 }
292
293 smol::future::yield_now().await;
294
295 Ok(())
296 }
297
298 /// Handles the stderr output from the context server.
299 /// Continuously reads and logs any error messages from the server.
300 async fn handle_err(transport: Arc<dyn Transport>) -> anyhow::Result<()> {
301 while let Some(err) = transport.receive_err().next().await {
302 log::warn!("context server stderr: {}", err.trim());
303 }
304
305 Ok(())
306 }
307
308 /// Handles the output to the context server's stdin.
309 /// This function continuously receives messages from the outbound channel,
310 /// writes them to the server's stdin, and manages the lifecycle of response handlers.
311 async fn handle_output(
312 transport: Arc<dyn Transport>,
313 outbound_rx: channel::Receiver<String>,
314 output_done_tx: barrier::Sender,
315 response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
316 ) -> anyhow::Result<()> {
317 let _clear_response_handlers = util::defer({
318 let response_handlers = response_handlers.clone();
319 move || {
320 response_handlers.lock().take();
321 }
322 });
323 while let Ok(message) = outbound_rx.recv().await {
324 log::trace!("outgoing message: {}", message);
325 transport.send(message).await?;
326 }
327 drop(output_done_tx);
328 Ok(())
329 }
330
331 /// Sends a JSON-RPC request to the context server and waits for a response.
332 /// This function handles serialization, deserialization, timeout, and error handling.
333 pub async fn request<T: DeserializeOwned>(
334 &self,
335 method: &str,
336 params: impl Serialize,
337 ) -> Result<T> {
338 self.request_impl(method, params, None).await
339 }
340
341 pub async fn cancellable_request<T: DeserializeOwned>(
342 &self,
343 method: &str,
344 params: impl Serialize,
345 cancel_rx: oneshot::Receiver<()>,
346 ) -> Result<T> {
347 self.request_impl(method, params, Some(cancel_rx)).await
348 }
349
350 pub async fn request_impl<T: DeserializeOwned>(
351 &self,
352 method: &str,
353 params: impl Serialize,
354 cancel_rx: Option<oneshot::Receiver<()>>,
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 = executor.timer(REQUEST_TIMEOUT).fuse();
391 let mut cancel_fut = pin!(
392 match cancel_rx {
393 Some(rx) => future::Either::Left(async {
394 rx.await.log_err();
395 }),
396 None => future::Either::Right(future::pending()),
397 }
398 .fuse()
399 );
400
401 select! {
402 response = rx.fuse() => {
403 let elapsed = started.elapsed();
404 log::trace!("took {elapsed:?} to receive response to {method:?} id {id}");
405 match response? {
406 Ok(response) => {
407 let parsed: AnyResponse = serde_json::from_str(&response)?;
408 if let Some(error) = parsed.error {
409 Err(anyhow!(error.message))
410 } else if let Some(result) = parsed.result {
411 Ok(serde_json::from_str(result.get())?)
412 } else {
413 anyhow::bail!("Invalid response: no result or error");
414 }
415 }
416 Err(_) => anyhow::bail!("cancelled")
417 }
418 }
419 _ = cancel_fut => {
420 self.notify(
421 Cancelled::METHOD,
422 ClientNotification::Cancelled(CancelledParams {
423 request_id: RequestId::Int(id),
424 reason: None
425 })
426 ).log_err();
427 anyhow::bail!("Request cancelled")
428 }
429 _ = timeout => {
430 log::error!("cancelled csp request task for {method:?} id {id} which took over {:?}", REQUEST_TIMEOUT);
431 anyhow::bail!("Context server request timeout");
432 }
433 }
434 }
435
436 /// Sends a notification to the context server without expecting a response.
437 /// This function serializes the notification and sends it through the outbound channel.
438 pub fn notify(&self, method: &str, params: impl Serialize) -> Result<()> {
439 let notification = serde_json::to_string(&Notification {
440 jsonrpc: JSON_RPC_VERSION,
441 method,
442 params,
443 })
444 .unwrap();
445 self.outbound_tx.try_send(notification)?;
446 Ok(())
447 }
448
449 #[allow(unused)]
450 pub fn on_notification<F>(&self, method: &'static str, f: F)
451 where
452 F: 'static + Send + FnMut(Value, AsyncApp),
453 {
454 self.notification_handlers
455 .lock()
456 .insert(method, Box::new(f));
457 }
458
459 pub fn on_request<R: crate::types::Request, F>(&self, mut f: F)
460 where
461 F: 'static + Send + FnMut(R::Params, AsyncApp) -> Task<Result<R::Response>>,
462 {
463 let outbound_tx = self.outbound_tx.clone();
464 self.request_handlers.lock().insert(
465 R::METHOD,
466 Box::new(move |id, json, cx| {
467 let outbound_tx = outbound_tx.clone();
468 match serde_json::from_str(json.get()) {
469 Ok(req) => {
470 let task = f(req, cx.clone());
471 cx.foreground_executor()
472 .spawn(async move {
473 match task.await {
474 Ok(res) => {
475 outbound_tx
476 .send(
477 serde_json::to_string(&Response {
478 jsonrpc: JSON_RPC_VERSION,
479 id,
480 value: CspResult::Ok(Some(res)),
481 })
482 .unwrap(),
483 )
484 .await
485 .ok();
486 }
487 Err(e) => {
488 outbound_tx
489 .send(
490 serde_json::to_string(&Response {
491 jsonrpc: JSON_RPC_VERSION,
492 id,
493 value: CspResult::<()>::Error(Some(Error {
494 code: -1, // todo!()
495 message: format!("{e}"),
496 })),
497 })
498 .unwrap(),
499 )
500 .await
501 .ok();
502 }
503 }
504 })
505 .detach();
506 }
507 Err(e) => {
508 cx.foreground_executor()
509 .spawn(async move {
510 outbound_tx
511 .send(
512 serde_json::to_string(&Response {
513 jsonrpc: JSON_RPC_VERSION,
514 id,
515 value: CspResult::<()>::Error(Some(Error {
516 code: -1, // todo!()
517 message: format!("{e}"),
518 })),
519 })
520 .unwrap(),
521 )
522 .await
523 .ok();
524 })
525 .detach();
526 }
527 }
528 }),
529 );
530 }
531}
532
533impl fmt::Display for ContextServerId {
534 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
535 self.0.fmt(f)
536 }
537}
538
539impl fmt::Debug for Client {
540 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
541 f.debug_struct("Context Server Client")
542 .field("id", &self.server_id.0)
543 .field("name", &self.name)
544 .finish_non_exhaustive()
545 }
546}