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