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