1use std::sync::Arc;
 2
 3use extension::{Extension, ExtensionContextServerProxy, ExtensionHostProxy, ProjectDelegate};
 4use gpui::{App, Entity};
 5
 6use crate::{ContextServerFactoryRegistry, ServerCommand};
 7
 8struct ExtensionProject {
 9    worktree_ids: Vec<u64>,
10}
11
12impl ProjectDelegate for ExtensionProject {
13    fn worktree_ids(&self) -> Vec<u64> {
14        self.worktree_ids.clone()
15    }
16}
17
18pub fn init(cx: &mut App) {
19    let proxy = ExtensionHostProxy::default_global(cx);
20    proxy.register_context_server_proxy(ContextServerFactoryRegistryProxy {
21        context_server_factory_registry: ContextServerFactoryRegistry::global(cx),
22    });
23}
24
25struct ContextServerFactoryRegistryProxy {
26    context_server_factory_registry: Entity<ContextServerFactoryRegistry>,
27}
28
29impl ExtensionContextServerProxy for ContextServerFactoryRegistryProxy {
30    fn register_context_server(&self, extension: Arc<dyn Extension>, id: Arc<str>, cx: &mut App) {
31        self.context_server_factory_registry
32            .update(cx, |registry, _| {
33                registry.register_server_factory(
34                    id.clone(),
35                    Arc::new({
36                        move |project, cx| {
37                            log::info!(
38                                "loading command for context server {id} from extension {}",
39                                extension.manifest().id
40                            );
41
42                            let id = id.clone();
43                            let extension = extension.clone();
44                            cx.spawn(|mut cx| async move {
45                                let extension_project =
46                                    project.update(&mut cx, |project, cx| {
47                                        Arc::new(ExtensionProject {
48                                            worktree_ids: project
49                                                .visible_worktrees(cx)
50                                                .map(|worktree| worktree.read(cx).id().to_proto())
51                                                .collect(),
52                                        })
53                                    })?;
54
55                                let command = extension
56                                    .context_server_command(id.clone(), extension_project)
57                                    .await?;
58
59                                log::info!("loaded command for context server {id}: {command:?}");
60
61                                Ok(ServerCommand {
62                                    path: command.command,
63                                    args: command.args,
64                                    env: Some(command.env.into_iter().collect()),
65                                })
66                            })
67                        }
68                    }),
69                )
70            });
71    }
72}