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