1use std::pin::Pin;
 2use std::sync::Arc;
 3
 4use anyhow::{anyhow, Result};
 5use async_trait::async_trait;
 6use context_servers::manager::{NativeContextServer, ServerConfig};
 7use context_servers::protocol::InitializedContextServerProtocol;
 8use context_servers::ContextServer;
 9use extension_host::wasm_host::{WasmExtension, WasmHost};
10use futures::{Future, FutureExt};
11use gpui::AsyncAppContext;
12
13pub struct ExtensionContextServer {
14    #[allow(unused)]
15    pub(crate) extension: WasmExtension,
16    #[allow(unused)]
17    pub(crate) host: Arc<WasmHost>,
18    id: Arc<str>,
19    context_server: Arc<NativeContextServer>,
20}
21
22impl ExtensionContextServer {
23    pub async fn new(extension: WasmExtension, host: Arc<WasmHost>, id: Arc<str>) -> Result<Self> {
24        let command = extension
25            .call({
26                let id = id.clone();
27                |extension, store| {
28                    async move {
29                        let command = extension
30                            .call_context_server_command(store, id.clone())
31                            .await?
32                            .map_err(|e| anyhow!("{}", e))?;
33                        anyhow::Ok(command)
34                    }
35                    .boxed()
36                }
37            })
38            .await?;
39
40        let config = Arc::new(ServerConfig {
41            id: id.to_string(),
42            executable: command.command,
43            args: command.args,
44            env: Some(command.env.into_iter().collect()),
45        });
46
47        anyhow::Ok(Self {
48            extension,
49            host,
50            id,
51            context_server: Arc::new(NativeContextServer::new(config)),
52        })
53    }
54}
55
56#[async_trait(?Send)]
57impl ContextServer for ExtensionContextServer {
58    fn id(&self) -> Arc<str> {
59        self.id.clone()
60    }
61
62    fn config(&self) -> Arc<ServerConfig> {
63        self.context_server.config()
64    }
65
66    fn client(&self) -> Option<Arc<InitializedContextServerProtocol>> {
67        self.context_server.client()
68    }
69
70    fn start<'a>(
71        self: Arc<Self>,
72        cx: &'a AsyncAppContext,
73    ) -> Pin<Box<dyn 'a + Future<Output = Result<()>>>> {
74        self.context_server.clone().start(cx)
75    }
76
77    fn stop(&self) -> Result<()> {
78        self.context_server.stop()
79    }
80}