tenant.openagents/omega

1use std::sync::Arc;
2
3use anyhow::Result;
4use context_server::ContextServerCommand;
5use extension::{
6    ContextServerConfiguration, Extension, ExtensionContextServerProxy, ExtensionHostProxy,
7    ProjectDelegate,
8};
9use gpui::{App, AsyncApp, Entity, Task};
10
11use crate::worktree_store::WorktreeStore;
12
13use super::registry::{self, ContextServerDescriptorRegistry};
14
15pub fn init(cx: &mut App) {
16    let proxy = ExtensionHostProxy::default_global(cx);
17    proxy.register_context_server_proxy(ContextServerDescriptorRegistryProxy {
18        context_server_factory_registry: ContextServerDescriptorRegistry::default_global(cx),
19    });
20}
21
22struct ExtensionProject {
23    worktree_ids: Vec<u64>,
24}
25
26impl ProjectDelegate for ExtensionProject {
27    fn worktree_ids(&self) -> Vec<u64> {
28        self.worktree_ids.clone()
29    }
30}
31
32struct ContextServerDescriptor {
33    id: Arc<str>,
34    extension: Arc<dyn Extension>,
35}
36
37fn extension_project(
38    worktree_store: Entity<WorktreeStore>,
39    cx: &mut AsyncApp,
40) -> Result<Arc<ExtensionProject>> {
41    Ok(worktree_store.update(cx, |worktree_store, cx| {
42        Arc::new(ExtensionProject {
43            worktree_ids: worktree_store
44                .visible_worktrees(cx)
45                .map(|worktree| worktree.read(cx).id().to_proto())
46                .collect(),
47        })
48    }))
49}
50
51impl registry::ContextServerDescriptor for ContextServerDescriptor {
52    fn command(
53        &self,
54        worktree_store: Entity<WorktreeStore>,
55        cx: &AsyncApp,
56    ) -> Task<Result<ContextServerCommand>> {
57        let id = self.id.clone();
58        let extension = self.extension.clone();
59        cx.spawn(async move |cx| {
60            let extension_project = extension_project(worktree_store, cx)?;
61            let mut command = extension
62                .context_server_command(id.clone(), extension_project.clone())
63                .await?;
64            command.command = extension.path_from_extension(&command.command);
65
66            log::debug!("loaded command for context server {id}: {command:?}");
67
68            Ok(ContextServerCommand {
69                path: command.command,
70                args: command.args,
71                env: Some(command.env.into_iter().collect()),
72                timeout: None,
73            })
74        })
75    }
76
77    fn configuration(
78        &self,
79        worktree_store: Entity<WorktreeStore>,
80        cx: &AsyncApp,
81    ) -> Task<Result<Option<ContextServerConfiguration>>> {
82        let id = self.id.clone();
83        let extension = self.extension.clone();
84        cx.spawn(async move |cx| {
85            let extension_project = extension_project(worktree_store, cx)?;
86            let configuration = extension
87                .context_server_configuration(id.clone(), extension_project)
88                .await?;
89
90            log::debug!("loaded configuration for context server {id}: {configuration:?}");
91
92            Ok(configuration)
93        })
94    }
95}
96
97struct ContextServerDescriptorRegistryProxy {
98    context_server_factory_registry: Entity<ContextServerDescriptorRegistry>,
99}
100
101impl ExtensionContextServerProxy for ContextServerDescriptorRegistryProxy {
102    fn register_context_server(&self, extension: Arc<dyn Extension>, id: Arc<str>, cx: &mut App) {
103        self.context_server_factory_registry
104            .update(cx, |registry, cx| {
105                registry.register_context_server_descriptor(
106                    id.clone(),
107                    Arc::new(ContextServerDescriptor { id, extension })
108                        as Arc<dyn registry::ContextServerDescriptor>,
109                    cx,
110                )
111            });
112    }
113
114    fn unregister_context_server(&self, server_id: Arc<str>, cx: &mut App) {
115        self.context_server_factory_registry
116            .update(cx, |registry, cx| {
117                registry.unregister_context_server_descriptor_by_id(&server_id, cx)
118            });
119    }
120}
121