use anyhow::{anyhow, Context as _, Result};use collections::HashMap;use futures::{channel::oneshot, io::BufWriter, select, AsyncRead, AsyncWrite, FutureExt};use gpui::{AppContext as _, AsyncApp, BackgroundExecutor, Task};use parking_lot::Mutex;use postage::barrier;use serde::{de::DeserializeOwned, Deserialize, Serialize};use serde_json::{value::RawValue, Value};use smol::{    channel,    io::{AsyncBufReadExt, AsyncWriteExt, BufReader},    process::Child,};use std::{    fmt,    path::PathBuf,    sync::{        atomic::{AtomicI32, Ordering::SeqCst},        Arc,    },    time::{Duration, Instant},};use util::TryFutureExt; const JSON_RPC_VERSION: &str = "2.0";const REQUEST_TIMEOUT: Duration = Duration::from_secs(60); // Standard JSON-RPC error codespub const PARSE_ERROR: i32 = -32700;pub const INVALID_REQUEST: i32 = -32600;pub const METHOD_NOT_FOUND: i32 = -32601;pub const INVALID_PARAMS: i32 = -32602;pub const INTERNAL_ERROR: i32 = -32603; type ResponseHandler = Box<dyn Send + FnOnce(Result<String, Error>)>;type NotificationHandler = Box<dyn Send + FnMut(Value, AsyncApp)>; #[derive(Debug, Clone, Eq, PartialEq, Hash, Serialize, Deserialize)]#[serde(untagged)]pub enum RequestId {    Int(i32),    Str(String),} pub struct Client {    server_id: ContextServerId,    next_id: AtomicI32,    outbound_tx: channel::Sender<String>,    name: Arc<str>,    notification_handlers: Arc<Mutex<HashMap<&'static str, NotificationHandler>>>,    response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,    #[allow(clippy::type_complexity)]    #[allow(dead_code)]    io_tasks: Mutex<Option<(Task<Option<()>>, Task<Option<()>>)>>,    #[allow(dead_code)]    output_done_rx: Mutex<Option<barrier::Receiver>>,    executor: BackgroundExecutor,    server: Arc<Mutex<Option<Child>>>,} #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]#[repr(transparent)]pub struct ContextServerId(pub Arc<str>); fn is_null_value<T: Serialize>(value: &T) -> bool {    if let Ok(Value::Null) = serde_json::to_value(value) {        true    } else {        false    }} #[derive(Serialize, Deserialize)]struct Request<'a, T> {    jsonrpc: &'static str,    id: RequestId,    method: &'a str,    #[serde(skip_serializing_if = "is_null_value")]    params: T,} #[derive(Serialize, Deserialize)]struct AnyResponse<'a> {    jsonrpc: &'a str,    id: RequestId,    #[serde(default)]    error: Option<Error>,    #[serde(borrow)]    result: Option<&'a RawValue>,} #[derive(Deserialize)]#[allow(dead_code)]struct Response<T> {    jsonrpc: &'static str,    id: RequestId,    #[serde(flatten)]    value: CspResult<T>,} #[derive(Deserialize)]#[serde(rename_all = "snake_case")]enum CspResult<T> {    #[serde(rename = "result")]    Ok(Option<T>),    #[allow(dead_code)]    Error(Option<Error>),} #[derive(Serialize, Deserialize)]struct Notification<'a, T> {    jsonrpc: &'static str,    #[serde(borrow)]    method: &'a str,    params: T,} #[derive(Debug, Clone, Deserialize)]struct AnyNotification<'a> {    jsonrpc: &'a str,    method: String,    #[serde(default)]    params: Option<Value>,} #[derive(Debug, Serialize, Deserialize)]struct Error {    message: String,} #[derive(Debug, Clone, Deserialize)]pub struct ModelContextServerBinary {    pub executable: PathBuf,    pub args: Vec<String>,    pub env: Option<HashMap<String, String>>,} impl Client {    /// Creates a new Client instance for a context server.    ///    /// This function initializes a new Client by spawning a child process for the context server,    /// setting up communication channels, and initializing handlers for input/output operations.    /// It takes a server ID, binary information, and an async app context as input.    pub fn new(        server_id: ContextServerId,        binary: ModelContextServerBinary,        cx: AsyncApp,    ) -> Result<Self> {        log::info!(            "starting context server (executable={:?}, args={:?})",            binary.executable,            &binary.args        );         let mut command = util::command::new_smol_command(&binary.executable);        command            .args(&binary.args)            .envs(binary.env.unwrap_or_default())            .stdin(std::process::Stdio::piped())            .stdout(std::process::Stdio::piped())            .stderr(std::process::Stdio::piped())            .kill_on_drop(true);         let mut server = command.spawn().with_context(|| {            format!(                "failed to spawn command. (path={:?}, args={:?})",                binary.executable, &binary.args            )        })?;         let stdin = server.stdin.take().unwrap();        let stdout = server.stdout.take().unwrap();        let stderr = server.stderr.take().unwrap();         let (outbound_tx, outbound_rx) = channel::unbounded::<String>();        let (output_done_tx, output_done_rx) = barrier::channel();         let notification_handlers =            Arc::new(Mutex::new(HashMap::<_, NotificationHandler>::default()));        let response_handlers =            Arc::new(Mutex::new(Some(HashMap::<_, ResponseHandler>::default())));         let stdout_input_task = cx.spawn({            let notification_handlers = notification_handlers.clone();            let response_handlers = response_handlers.clone();            move |cx| {                Self::handle_input(stdout, notification_handlers, response_handlers, cx).log_err()            }        });        let stderr_input_task = cx.spawn(|_| Self::handle_stderr(stderr).log_err());        let input_task = cx.spawn(|_| async move {            let (stdout, stderr) = futures::join!(stdout_input_task, stderr_input_task);            stdout.or(stderr)        });        let output_task = cx.background_spawn({            Self::handle_output(                stdin,                outbound_rx,                output_done_tx,                response_handlers.clone(),            )            .log_err()        });         let mut context_server = Self {            server_id,            notification_handlers,            response_handlers,            name: "".into(),            next_id: Default::default(),            outbound_tx,            executor: cx.background_executor().clone(),            io_tasks: Mutex::new(Some((input_task, output_task))),            output_done_rx: Mutex::new(Some(output_done_rx)),            server: Arc::new(Mutex::new(Some(server))),        };         if let Some(name) = binary.executable.file_name() {            context_server.name = name.to_string_lossy().into();        }         Ok(context_server)    }     /// Handles input from the server's stdout.    ///    /// This function continuously reads lines from the provided stdout stream,    /// parses them as JSON-RPC responses or notifications, and dispatches them    /// to the appropriate handlers. It processes both responses (which are matched    /// to pending requests) and notifications (which trigger registered handlers).    async fn handle_input<Stdout>(        stdout: Stdout,        notification_handlers: Arc<Mutex<HashMap<&'static str, NotificationHandler>>>,        response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,        cx: AsyncApp,    ) -> anyhow::Result<()>    where        Stdout: AsyncRead + Unpin + Send + 'static,    {        let mut stdout = BufReader::new(stdout);        let mut buffer = String::new();         loop {            buffer.clear();            if stdout.read_line(&mut buffer).await? == 0 {                return Ok(());            }             let content = buffer.trim();             if !content.is_empty() {                if let Ok(response) = serde_json::from_str::<AnyResponse>(content) {                    if let Some(handlers) = response_handlers.lock().as_mut() {                        if let Some(handler) = handlers.remove(&response.id) {                            handler(Ok(content.to_string()));                        }                    }                } else if let Ok(notification) = serde_json::from_str::<AnyNotification>(content) {                    let mut notification_handlers = notification_handlers.lock();                    if let Some(handler) =                        notification_handlers.get_mut(notification.method.as_str())                    {                        handler(notification.params.unwrap_or(Value::Null), cx.clone());                    }                }            }             smol::future::yield_now().await;        }    }     /// Handles the stderr output from the context server.    /// Continuously reads and logs any error messages from the server.    async fn handle_stderr<Stderr>(stderr: Stderr) -> anyhow::Result<()>    where        Stderr: AsyncRead + Unpin + Send + 'static,    {        let mut stderr = BufReader::new(stderr);        let mut buffer = String::new();         loop {            buffer.clear();            if stderr.read_line(&mut buffer).await? == 0 {                return Ok(());            }            log::warn!("context server stderr: {}", buffer.trim());            smol::future::yield_now().await;        }    }     /// Handles the output to the context server's stdin.    /// This function continuously receives messages from the outbound channel,    /// writes them to the server's stdin, and manages the lifecycle of response handlers.    async fn handle_output<Stdin>(        stdin: Stdin,        outbound_rx: channel::Receiver<String>,        output_done_tx: barrier::Sender,        response_handlers: Arc<Mutex<Option<HashMap<RequestId, ResponseHandler>>>>,    ) -> anyhow::Result<()>    where        Stdin: AsyncWrite + Unpin + Send + 'static,    {        let mut stdin = BufWriter::new(stdin);        let _clear_response_handlers = util::defer({            let response_handlers = response_handlers.clone();            move || {                response_handlers.lock().take();            }        });        while let Ok(message) = outbound_rx.recv().await {            log::trace!("outgoing message: {}", message);             stdin.write_all(message.as_bytes()).await?;            stdin.write_all(b"\n").await?;            stdin.flush().await?;        }        drop(output_done_tx);        Ok(())    }     /// Sends a JSON-RPC request to the context server and waits for a response.    /// This function handles serialization, deserialization, timeout, and error handling.    pub async fn request<T: DeserializeOwned>(        &self,        method: &str,        params: impl Serialize,    ) -> Result<T> {        let id = self.next_id.fetch_add(1, SeqCst);        let request = serde_json::to_string(&Request {            jsonrpc: JSON_RPC_VERSION,            id: RequestId::Int(id),            method,            params,        })        .unwrap();         let (tx, rx) = oneshot::channel();        let handle_response = self            .response_handlers            .lock()            .as_mut()            .ok_or_else(|| anyhow!("server shut down"))            .map(|handlers| {                handlers.insert(                    RequestId::Int(id),                    Box::new(move |result| {                        let _ = tx.send(result);                    }),                );            });         let send = self            .outbound_tx            .try_send(request)            .context("failed to write to context server's stdin");         let executor = self.executor.clone();        let started = Instant::now();        handle_response?;        send?;         let mut timeout = executor.timer(REQUEST_TIMEOUT).fuse();        select! {            response = rx.fuse() => {                let elapsed = started.elapsed();                log::trace!("took {elapsed:?} to receive response to {method:?} id {id}");                match response? {                    Ok(response) => {                        let parsed: AnyResponse = serde_json::from_str(&response)?;                        if let Some(error) = parsed.error {                            Err(anyhow!(error.message))                        } else if let Some(result) = parsed.result {                            Ok(serde_json::from_str(result.get())?)                        } else {                            Err(anyhow!("Invalid response: no result or error"))                        }                    }                    Err(_) => anyhow::bail!("cancelled")                }            }            _ = timeout => {                log::error!("cancelled csp request task for {method:?} id {id} which took over {:?}", REQUEST_TIMEOUT);                anyhow::bail!("Context server request timeout");            }        }    }     /// Sends a notification to the context server without expecting a response.    /// This function serializes the notification and sends it through the outbound channel.    pub fn notify(&self, method: &str, params: impl Serialize) -> Result<()> {        let notification = serde_json::to_string(&Notification {            jsonrpc: JSON_RPC_VERSION,            method,            params,        })        .unwrap();        self.outbound_tx.try_send(notification)?;        Ok(())    }     pub fn on_notification<F>(&self, method: &'static str, f: F)    where        F: 'static + Send + FnMut(Value, AsyncApp),    {        self.notification_handlers            .lock()            .insert(method, Box::new(f));    }     pub fn name(&self) -> &str {        &self.name    }     pub fn server_id(&self) -> ContextServerId {        self.server_id.clone()    }} impl Drop for Client {    fn drop(&mut self) {        if let Some(mut server) = self.server.lock().take() {            let _ = server.kill();        }    }} impl fmt::Display for ContextServerId {    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {        self.0.fmt(f)    }} impl fmt::Debug for Client {    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {        f.debug_struct("Context Server Client")            .field("id", &self.server_id.0)            .field("name", &self.name)            .finish_non_exhaustive()    }}