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