tenant.openagents/omega

1use anyhow::{Context as _, Result, anyhow};
2use collections::HashMap;
3use futures::{FutureExt, StreamExt, channel::oneshot, future, select};
4use futures_lite::future::yield_now;
5use gpui::{AppContext as _, AsyncApp, BackgroundExecutor, Task};
6use parking_lot::Mutex;
7use postage::{barrier, prelude::Stream as _};
8use serde::{Deserialize, Serialize, de::DeserializeOwned};
9use serde_json::{Value, value::RawValue};
10use slotmap::SlotMap;
11use std::{
12    fmt,
13    path::PathBuf,
14    pin::pin,
15    sync::{
16        Arc,
17        atomic::{AtomicI32, Ordering::SeqCst},
18    },
19    time::{Duration, Instant},
20};
21use util::{ResultExt, TryFutureExt};
22
23use crate::{
24    oauth::WwwAuthenticate,
25    transport::{StdioTransport, Transport},
26    types::{CancelledParams, ClientNotification, Notification as _, notifications::Cancelled},
27};
28
29const JSON_RPC_VERSION: &str = "2.0";
30const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
31
32// Standard JSON-RPC error codes
33pub const PARSE_ERROR: i32 = -32700;
34pub const INVALID_REQUEST: i32 = -32600;
35pub const METHOD_NOT_FOUND: i32 = -32601;
36pub const INVALID_PARAMS: i32 = -32602;
37pub const INTERNAL_ERROR: i32 = -32603;
38
39type ResponseHandler = Box<dyn Send + FnOnce(String)>;
40type NotificationHandler = Box<dyn Send + FnMut(Value, AsyncApp)>;
41type RequestHandler = Box<dyn Send + FnMut(RequestId, &RawValue, AsyncApp)>;
42
43#[derive(Debug, Clone, Eq, PartialEq, Hash, Serialize, Deserialize)]
44#[serde(untagged)]
45pub enum RequestId {
46    Int(i32),
47    Str(String),
48}
49
50pub(crate) struct Client {
51    server_id: ContextServerId,
52    next_id: AtomicI32,
53    outbound_tx: async_channel::Sender<String>,
54    name: Arc<str>,
55    subscription_set: Arc<Mutex<NotificationSubscriptionSet>>,
56    response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
57    #[allow(clippy::type_complexity)]
58    #[allow(dead_code)]
59    io_tasks: Mutex<Option<(Task<Option<()>>, Task<Option<()>>)>>,
60    output_done_rx: Mutex<Option<barrier::Receiver>>,
61    executor: BackgroundExecutor,
62    transport: Arc<dyn Transport>,
63    request_timeout: Option<Duration>,
64    /// Single-slot side channel for the last transport-level error. When the
65    /// output task encounters a send failure it stashes the error here and
66    /// exits; the next request to observe cancellation `.take()`s it so it
67    /// can fail with the underlying cause (e.g. "connection refused") instead
68    /// of a generic "cancelled". This is best-effort diagnostics: with
69    /// concurrent requests in flight, a single arbitrary one receives the
70    /// stashed error. Nothing may depend on it for correctness —
71    /// authentication challenges are observed via [`Self::wait_for_shutdown`]
72    /// and [`Transport::auth_challenge`], which do not involve requests.
73    last_transport_error: Arc<Mutex<Option<anyhow::Error>>>,
74}
75
76#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
77#[repr(transparent)]
78pub(crate) struct ContextServerId(pub Arc<str>);
79
80fn is_null_value<T: Serialize>(value: &T) -> bool {
81    matches!(serde_json::to_value(value), Ok(Value::Null))
82}
83
84#[derive(Serialize, Deserialize)]
85pub struct Request<'a, T> {
86    pub jsonrpc: &'static str,
87    pub id: RequestId,
88    pub method: &'a str,
89    #[serde(skip_serializing_if = "is_null_value")]
90    pub params: T,
91}
92
93#[derive(Serialize, Deserialize)]
94pub struct AnyRequest<'a> {
95    pub jsonrpc: &'a str,
96    pub id: RequestId,
97    pub method: &'a str,
98    #[serde(skip_serializing_if = "is_null_value")]
99    pub params: Option<&'a RawValue>,
100}
101
102#[derive(Serialize, Deserialize)]
103struct AnyResponse<'a> {
104    jsonrpc: &'a str,
105    id: RequestId,
106    #[serde(default)]
107    error: Option<Error>,
108    #[serde(borrow)]
109    result: Option<&'a RawValue>,
110}
111
112#[derive(Serialize, Deserialize)]
113#[allow(dead_code)]
114pub(crate) struct Response<T> {
115    pub jsonrpc: &'static str,
116    pub id: RequestId,
117    #[serde(flatten)]
118    pub value: CspResult<T>,
119}
120
121#[derive(Serialize, Deserialize)]
122#[serde(rename_all = "snake_case")]
123pub(crate) enum CspResult<T> {
124    #[serde(rename = "result")]
125    Ok(Option<T>),
126    #[allow(dead_code)]
127    Error(Option<Error>),
128}
129
130#[derive(Serialize, Deserialize)]
131struct Notification<'a, T> {
132    jsonrpc: &'static str,
133    #[serde(borrow)]
134    method: &'a str,
135    #[serde(skip_serializing_if = "is_null_value")]
136    params: T,
137}
138
139#[derive(Debug, Clone, Deserialize)]
140struct AnyNotification<'a> {
141    #[expect(
142        unused,
143        reason = "Part of the JSON-RPC protocol - we expect the field to be present in a valid JSON-RPC notification"
144    )]
145    jsonrpc: &'a str,
146    method: String,
147    #[serde(default)]
148    params: Option<Value>,
149}
150
151#[derive(Debug, Serialize, Deserialize)]
152pub(crate) struct Error {
153    pub message: String,
154    pub code: i32,
155}
156
157#[derive(Debug, Clone, Deserialize)]
158pub struct ModelContextServerBinary {
159    pub executable: PathBuf,
160    pub args: Vec<String>,
161    pub env: Option<HashMap<String, String>>,
162    pub timeout: Option<u64>,
163}
164
165impl Client {
166    /// Creates a new Client instance for a context server.
167    ///
168    /// This function initializes a new Client by spawning a child process for the context server,
169    /// setting up communication channels, and initializing handlers for input/output operations.
170    /// It takes a server ID, binary information, and an async app context as input.
171    pub fn stdio(
172        server_id: ContextServerId,
173        binary: ModelContextServerBinary,
174        working_directory: &Option<PathBuf>,
175        cx: AsyncApp,
176    ) -> Result<Self> {
177        log::debug!(
178            "starting context server (executable={:?}, args={:?})",
179            binary.executable,
180            &binary.args
181        );
182
183        let server_name = binary
184            .executable
185            .file_name()
186            .map(|name| name.to_string_lossy().into_owned())
187            .unwrap_or_else(String::new);
188
189        let timeout = binary.timeout.map(Duration::from_secs);
190        let transport = Arc::new(StdioTransport::new(binary, working_directory, &cx)?);
191        Self::new(server_id, server_name.into(), transport, timeout, cx)
192    }
193
194    /// Creates a new Client instance for a context server.
195    pub fn new(
196        server_id: ContextServerId,
197        server_name: Arc<str>,
198        transport: Arc<dyn Transport>,
199        request_timeout: Option<Duration>,
200        cx: AsyncApp,
201    ) -> Result<Self> {
202        let (outbound_tx, outbound_rx) = async_channel::unbounded::<String>();
203        let (output_done_tx, output_done_rx) = barrier::channel();
204
205        let subscription_set = Arc::new(Mutex::new(NotificationSubscriptionSet::default()));
206        let response_handlers =
207            Arc::new(Mutex::new(Some(HashMap::<_, ResponseHandler>::default())));
208        let request_handlers = Arc::new(Mutex::new(HashMap::<_, RequestHandler>::default()));
209
210        let receive_input_task = cx.spawn({
211            let subscription_set = subscription_set.clone();
212            let response_handlers = response_handlers.clone();
213            let request_handlers = request_handlers.clone();
214            let transport = transport.clone();
215            async move |cx| {
216                Self::handle_input(
217                    transport,
218                    subscription_set,
219                    request_handlers,
220                    response_handlers,
221                    cx,
222                )
223                .log_err()
224                .await
225            }
226        });
227        let receive_err_task = cx.spawn({
228            let transport = transport.clone();
229            async move |_| Self::handle_err(transport).log_err().await
230        });
231        let input_task = cx.spawn(async move |_| {
232            let (input, err) = futures::join!(receive_input_task, receive_err_task);
233            input.or(err)
234        });
235
236        let last_transport_error: Arc<Mutex<Option<anyhow::Error>>> = Arc::new(Mutex::new(None));
237        let output_task = cx.background_spawn({
238            let transport = transport.clone();
239            let last_transport_error = last_transport_error.clone();
240            Self::handle_output(
241                transport,
242                outbound_rx,
243                output_done_tx,
244                response_handlers.clone(),
245                last_transport_error,
246            )
247            .log_err()
248        });
249
250        Ok(Self {
251            server_id,
252            subscription_set,
253            response_handlers,
254            name: server_name,
255            next_id: Default::default(),
256            outbound_tx,
257            executor: cx.background_executor().clone(),
258            io_tasks: Mutex::new(Some((input_task, output_task))),
259            output_done_rx: Mutex::new(Some(output_done_rx)),
260            transport,
261            request_timeout,
262            last_transport_error,
263        })
264    }
265
266    /// Handles input from the server's stdout.
267    ///
268    /// This function continuously reads lines from the provided stdout stream,
269    /// parses them as JSON-RPC responses or notifications, and dispatches them
270    /// to the appropriate handlers. It processes both responses (which are matched
271    /// to pending requests) and notifications (which trigger registered handlers).
272    async fn handle_input(
273        transport: Arc<dyn Transport>,
274        subscription_set: Arc<Mutex<NotificationSubscriptionSet>>,
275        request_handlers: Arc<Mutex<HashMap<&'static str, RequestHandler>>>,
276        response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
277        cx: &mut AsyncApp,
278    ) -> anyhow::Result<()> {
279        let mut receiver = transport.receive();
280
281        while let Some(message) = receiver.next().await {
282            log::trace!("recv: {}", &message);
283            if let Ok(request) = serde_json::from_str::<AnyRequest>(&message) {
284                let mut request_handlers = request_handlers.lock();
285                if let Some(handler) = request_handlers.get_mut(request.method) {
286                    handler(
287                        request.id,
288                        request.params.unwrap_or(RawValue::NULL),
289                        cx.clone(),
290                    );
291                }
292            } else if let Ok(response) = serde_json::from_str::<AnyResponse>(&message) {
293                if let Some(handlers) = response_handlers.lock().as_mut()
294                    && let Some(handler) = handlers.remove(&response.id)
295                {
296                    handler(message.to_string());
297                }
298            } else if let Ok(notification) = serde_json::from_str::<AnyNotification>(&message) {
299                subscription_set.lock().notify(
300                    &notification.method,
301                    notification.params.unwrap_or(Value::Null),
302                    cx,
303                )
304            } else {
305                log::error!("Unhandled JSON from context_server: {}", message);
306            }
307        }
308
309        yield_now().await;
310
311        Ok(())
312    }
313
314    /// Handles the stderr output from the context server.
315    /// Continuously reads and logs any error messages from the server.
316    async fn handle_err(transport: Arc<dyn Transport>) -> anyhow::Result<()> {
317        while let Some(err) = transport.receive_err().next().await {
318            log::debug!("context server stderr: {}", err.trim());
319        }
320
321        Ok(())
322    }
323
324    /// Handles the output to the context server's stdin.
325    /// This function continuously receives messages from the outbound channel,
326    /// writes them to the server's stdin, and manages the lifecycle of response handlers.
327    async fn handle_output(
328        transport: Arc<dyn Transport>,
329        outbound_rx: async_channel::Receiver<String>,
330        output_done_tx: barrier::Sender,
331        response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,
332        last_transport_error: Arc<Mutex<Option<anyhow::Error>>>,
333    ) -> anyhow::Result<()> {
334        let _clear_response_handlers = util::defer({
335            let response_handlers = response_handlers.clone();
336            move || {
337                response_handlers.lock().take();
338            }
339        });
340        while let Ok(message) = outbound_rx.recv().await {
341            log::trace!("outgoing message: {}", message);
342            if let Err(err) = transport.send(message).await {
343                log::debug!("transport send failed: {:#}", err);
344                *last_transport_error.lock() = Some(err);
345                return Ok(());
346            }
347        }
348        drop(output_done_tx);
349        Ok(())
350    }
351
352    /// A future that resolves once the transport's output loop has terminated
353    /// — after a send failure, or when this client is dropped — yielding the
354    /// authentication challenge recorded by the transport if it shut down on a
355    /// `401 Unauthorized` response.
356    ///
357    /// Unlike `last_transport_error`, this does not require a request to be in
358    /// flight when the transport fails. Returns `None` if the shutdown signal
359    /// was already claimed: there is a single signal per client.
360    pub(crate) fn wait_for_shutdown(
361        &self,
362    ) -> Option<future::BoxFuture<'static, Option<WwwAuthenticate>>> {
363        let mut output_done = self.output_done_rx.lock().take()?;
364        let transport = self.transport.clone();
365        Some(
366            async move {
367                output_done.recv().await;
368                transport.auth_challenge()
369            }
370            .boxed(),
371        )
372    }
373
374    /// Sends a JSON-RPC request to the context server and waits for a response.
375    /// This function handles serialization, deserialization, timeout, and error handling.
376    pub async fn request<T: DeserializeOwned>(
377        &self,
378        method: &str,
379        params: impl Serialize,
380    ) -> Result<T> {
381        self.request_with(
382            method,
383            params,
384            None,
385            self.request_timeout.or(Some(DEFAULT_REQUEST_TIMEOUT)),
386        )
387        .await
388    }
389
390    pub async fn request_with<T: DeserializeOwned>(
391        &self,
392        method: &str,
393        params: impl Serialize,
394        cancel_rx: Option<oneshot::Receiver<()>>,
395        timeout: Option<Duration>,
396    ) -> Result<T> {
397        let id = self.next_id.fetch_add(1, SeqCst);
398        let request = serde_json::to_string(&Request {
399            jsonrpc: JSON_RPC_VERSION,
400            id: RequestId::Int(id),
401            method,
402            params,
403        })
404        .unwrap();
405
406        let (tx, rx) = oneshot::channel();
407        let handle_response = self
408            .response_handlers
409            .lock()
410            .as_mut()
411            .context("server shut down")
412            .map(|handlers| {
413                handlers.insert(
414                    RequestId::Int(id),
415                    Box::new(move |result| {
416                        let _ = tx.send(result);
417                    }),
418                );
419            });
420
421        let send = self
422            .outbound_tx
423            .try_send(request)
424            .context("failed to write to context server's stdin");
425
426        let executor = self.executor.clone();
427        let started = Instant::now();
428        handle_response?;
429        send?;
430
431        let mut timeout_fut = pin!(
432            match timeout {
433                Some(timeout) => future::Either::Left(executor.timer(timeout)),
434                None => future::Either::Right(future::pending()),
435            }
436            .fuse()
437        );
438        let mut cancel_fut = pin!(
439            match cancel_rx {
440                Some(rx) => future::Either::Left(async {
441                    rx.await.log_err();
442                }),
443                None => future::Either::Right(future::pending()),
444            }
445            .fuse()
446        );
447
448        select! {
449            response = rx.fuse() => {
450                let elapsed = started.elapsed();
451                log::trace!("took {elapsed:?} to receive response to {method:?} id {id}");
452                match response {
453                    Ok(response) => {
454                        let parsed: AnyResponse = serde_json::from_str(&response)?;
455                        if let Some(error) = parsed.error {
456                            Err(anyhow!(error.message))
457                        } else if let Some(result) = parsed.result {
458                            Ok(serde_json::from_str(result.get())?)
459                        } else {
460                            anyhow::bail!("Invalid response: no result or error");
461                        }
462                    }
463                    Err(_canceled) => {
464                        if let Some(err) = self.last_transport_error.lock().take() {
465                            return Err(err);
466                        }
467                        anyhow::bail!("cancelled")
468                    }
469                }
470            }
471            _ = cancel_fut => {
472                self.notify(
473                    Cancelled::METHOD,
474                    ClientNotification::Cancelled(CancelledParams {
475                        request_id: RequestId::Int(id),
476                        reason: None
477                    })
478                ).log_err();
479                anyhow::bail!(RequestCanceled)
480            }
481            _ = timeout_fut => {
482                log::error!("cancelled csp request task for {method:?} id {id} which took over {:?}", timeout.unwrap());
483                anyhow::bail!("Context server request timeout");
484            }
485        }
486    }
487
488    /// Sends a notification to the context server without expecting a response.
489    /// This function serializes the notification and sends it through the outbound channel.
490    pub fn notify(&self, method: &str, params: impl Serialize) -> Result<()> {
491        let notification = serde_json::to_string(&Notification {
492            jsonrpc: JSON_RPC_VERSION,
493            method,
494            params,
495        })
496        .unwrap();
497        self.outbound_tx.try_send(notification)?;
498        Ok(())
499    }
500
501    /// Notify the underlying transport of the negotiated MCP protocol version
502    /// so it can stamp subsequent requests (e.g. HTTP's `MCP-Protocol-Version`
503    /// header required from 2025-06-18 onward).
504    pub(crate) fn set_protocol_version(&self, version: &str) {
505        self.transport.set_protocol_version(version);
506    }
507
508    #[must_use]
509    pub fn on_notification(
510        &self,
511        method: &'static str,
512        f: Box<dyn 'static + Send + FnMut(Value, AsyncApp)>,
513    ) -> NotificationSubscription {
514        let mut notification_subscriptions = self.subscription_set.lock();
515
516        NotificationSubscription {
517            id: notification_subscriptions.add_handler(method, f),
518            set: self.subscription_set.clone(),
519        }
520    }
521}
522
523#[derive(Debug)]
524pub struct RequestCanceled;
525
526impl std::error::Error for RequestCanceled {}
527
528impl std::fmt::Display for RequestCanceled {
529    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
530        f.write_str("Context server request was canceled")
531    }
532}
533
534impl fmt::Display for ContextServerId {
535    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
536        self.0.fmt(f)
537    }
538}
539
540impl fmt::Debug for Client {
541    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
542        f.debug_struct("Context Server Client")
543            .field("id", &self.server_id.0)
544            .field("name", &self.name)
545            .finish_non_exhaustive()
546    }
547}
548
549slotmap::new_key_type! {
550    struct NotificationSubscriptionId;
551}
552
553#[derive(Default)]
554pub struct NotificationSubscriptionSet {
555    // we have very few subscriptions at the moment
556    methods: Vec<(&'static str, Vec<NotificationSubscriptionId>)>,
557    handlers: SlotMap<NotificationSubscriptionId, NotificationHandler>,
558}
559
560impl NotificationSubscriptionSet {
561    #[must_use]
562    fn add_handler(
563        &mut self,
564        method: &'static str,
565        handler: NotificationHandler,
566    ) -> NotificationSubscriptionId {
567        let id = self.handlers.insert(handler);
568        if let Some((_, handler_ids)) = self
569            .methods
570            .iter_mut()
571            .find(|(probe_method, _)| method == *probe_method)
572        {
573            debug_assert!(
574                handler_ids.len() < 20,
575                "Too many MCP handlers for {}. Consider using a different data structure.",
576                method
577            );
578
579            handler_ids.push(id);
580        } else {
581            self.methods.push((method, vec![id]));
582        };
583        id
584    }
585
586    fn notify(&mut self, method: &str, payload: Value, cx: &mut AsyncApp) {
587        let Some((_, handler_ids)) = self
588            .methods
589            .iter_mut()
590            .find(|(probe_method, _)| method == *probe_method)
591        else {
592            return;
593        };
594
595        for handler_id in handler_ids {
596            if let Some(handler) = self.handlers.get_mut(*handler_id) {
597                handler(payload.clone(), cx.clone());
598            }
599        }
600    }
601}
602
603pub struct NotificationSubscription {
604    id: NotificationSubscriptionId,
605    set: Arc<Mutex<NotificationSubscriptionSet>>,
606}
607
608impl Drop for NotificationSubscription {
609    fn drop(&mut self) {
610        let mut set = self.set.lock();
611        set.handlers.remove(self.id);
612        set.methods.retain_mut(|(_, handler_ids)| {
613            handler_ids.retain(|id| *id != self.id);
614            !handler_ids.is_empty()
615        });
616    }
617}
618