tenant.openagents/omega

1use std::rc::Rc;
2
3use acp_thread::{AgentConnection, LoadError};
4use agent_servers::AcpConnection;
5use agent_servers::{AgentServer, AgentServerDelegate};
6use anyhow::Result;
7use collections::HashMap;
8use futures::{FutureExt, future::Shared};
9use gpui::{App, AppContext, Context, Entity, EventEmitter, SharedString, Subscription, Task};
10
11use project::{AgentServerStore, AgentServersUpdated, Project};
12use watch::Receiver;
13
14use crate::Agent;
15
16pub enum AgentConnectionEntry {
17    Connecting {
18        connect_task: Shared<Task<Result<AgentConnectedState, LoadError>>>,
19    },
20    Connected(AgentConnectedState),
21    Error {
22        error: LoadError,
23    },
24}
25
26#[derive(Clone)]
27pub struct AgentConnectedState {
28    pub connection: Rc<dyn AgentConnection>,
29}
30
31#[derive(Clone, Copy, Debug, PartialEq, Eq)]
32pub enum AgentConnectionStatus {
33    Disconnected,
34    Connecting,
35    Connected,
36}
37
38impl AgentConnectionEntry {
39    pub fn wait_for_connection(&self) -> Shared<Task<Result<AgentConnectedState, LoadError>>> {
40        match self {
41            AgentConnectionEntry::Connecting { connect_task } => connect_task.clone(),
42            AgentConnectionEntry::Connected(state) => Task::ready(Ok(state.clone())).shared(),
43            AgentConnectionEntry::Error { error } => Task::ready(Err(error.clone())).shared(),
44        }
45    }
46
47    pub fn status(&self) -> AgentConnectionStatus {
48        match self {
49            AgentConnectionEntry::Connecting { .. } => AgentConnectionStatus::Connecting,
50            AgentConnectionEntry::Connected(_) => AgentConnectionStatus::Connected,
51            AgentConnectionEntry::Error { .. } => AgentConnectionStatus::Disconnected,
52        }
53    }
54}
55
56pub enum AgentConnectionEntryEvent {
57    NewVersionAvailable(SharedString),
58    LoadingStatusChanged(Option<SharedString>),
59}
60
61impl EventEmitter<AgentConnectionEntryEvent> for AgentConnectionEntry {}
62
63#[derive(Clone)]
64pub struct ActiveAcpConnection {
65    pub agent_id: project::AgentId,
66    pub connection: Rc<AcpConnection>,
67}
68
69pub struct AgentConnectionStore {
70    project: Entity<Project>,
71    entries: HashMap<Agent, Entity<AgentConnectionEntry>>,
72    _subscriptions: Vec<Subscription>,
73}
74
75impl AgentConnectionStore {
76    pub fn new(project: Entity<Project>, cx: &mut Context<Self>) -> Self {
77        let agent_server_store = project.read(cx).agent_server_store().clone();
78        let subscription = cx.subscribe(&agent_server_store, Self::handle_agent_servers_updated);
79        Self {
80            project,
81            entries: HashMap::default(),
82            _subscriptions: vec![subscription],
83        }
84    }
85
86    pub fn project(&self) -> &Entity<Project> {
87        &self.project
88    }
89
90    pub fn entry(&self, key: &Agent) -> Option<&Entity<AgentConnectionEntry>> {
91        self.entries.get(key)
92    }
93
94    pub fn connection_status(&self, key: &Agent, cx: &App) -> AgentConnectionStatus {
95        self.entries
96            .get(key)
97            .map(|entry| entry.read(cx).status())
98            .unwrap_or(AgentConnectionStatus::Disconnected)
99    }
100
101    pub fn agent_version(&self, key: &Agent, cx: &App) -> Option<SharedString> {
102        match self.entries.get(key)?.read(cx) {
103            AgentConnectionEntry::Connected(state) => state.connection.agent_version(),
104            AgentConnectionEntry::Connecting { .. } | AgentConnectionEntry::Error { .. } => None,
105        }
106    }
107
108    pub fn active_acp_connections(&self, cx: &App) -> Vec<ActiveAcpConnection> {
109        self.entries
110            .values()
111            .filter_map(|entry| match entry.read(cx) {
112                AgentConnectionEntry::Connected(state) => state
113                    .connection
114                    .clone()
115                    .downcast::<AcpConnection>()
116                    .map(|connection| ActiveAcpConnection {
117                        agent_id: state.connection.agent_id(),
118                        connection,
119                    }),
120                AgentConnectionEntry::Connecting { .. } | AgentConnectionEntry::Error { .. } => {
121                    None
122                }
123            })
124            .collect()
125    }
126
127    pub fn restart_connection(
128        &mut self,
129        key: Agent,
130        server: Rc<dyn AgentServer>,
131        cx: &mut Context<Self>,
132    ) -> Entity<AgentConnectionEntry> {
133        if let Some(entry) = self.entries.get(&key) {
134            if matches!(entry.read(cx), AgentConnectionEntry::Connecting { .. }) {
135                return entry.clone();
136            }
137        }
138
139        self.entries.remove(&key);
140        self.request_connection(key, server, cx)
141    }
142
143    pub fn request_connection(
144        &mut self,
145        key: Agent,
146        server: Rc<dyn AgentServer>,
147        cx: &mut Context<Self>,
148    ) -> Entity<AgentConnectionEntry> {
149        if let Some(entry) = self.entries.get(&key) {
150            return entry.clone();
151        }
152
153        let (mut new_version_rx, mut loading_status_rx, connect_task) =
154            self.start_connection(server, cx);
155        let connect_task = connect_task.shared();
156
157        let entry = cx.new(|_cx| AgentConnectionEntry::Connecting {
158            connect_task: connect_task.clone(),
159        });
160
161        self.entries.insert(key.clone(), entry.clone());
162        cx.notify();
163
164        cx.spawn({
165            let key = key.clone();
166            let entry = entry.downgrade();
167            async move |this, cx| match connect_task.await {
168                Ok(connected_state) => {
169                    this.update(cx, move |this, cx| {
170                        if this.entries.get(&key) != entry.upgrade().as_ref() {
171                            return;
172                        }
173
174                        entry
175                            .update(cx, move |entry, cx| {
176                                if let AgentConnectionEntry::Connecting { .. } = entry {
177                                    *entry = AgentConnectionEntry::Connected(connected_state);
178                                    cx.notify();
179                                }
180                            })
181                            .ok();
182                        cx.notify();
183                    })
184                    .ok();
185                }
186                Err(error) => {
187                    this.update(cx, move |this, cx| {
188                        if this.entries.get(&key) != entry.upgrade().as_ref() {
189                            return;
190                        }
191
192                        entry
193                            .update(cx, move |entry, cx| {
194                                if let AgentConnectionEntry::Connecting { .. } = entry {
195                                    *entry = AgentConnectionEntry::Error { error };
196                                    cx.notify();
197                                }
198                            })
199                            .ok();
200                        this.entries.remove(&key);
201                        cx.notify();
202                    })
203                    .ok();
204                }
205            }
206        })
207        .detach();
208
209        cx.spawn({
210            let key = key.clone();
211            let entry = entry.downgrade();
212            async move |this, cx| {
213                while let Ok(version) = new_version_rx.recv().await {
214                    let Some(version) = version else {
215                        continue;
216                    };
217
218                    this.update(cx, move |this, cx| {
219                        if this.entries.get(&key) != entry.upgrade().as_ref() {
220                            return;
221                        }
222
223                        entry
224                            .update(cx, move |_entry, cx| {
225                                cx.emit(AgentConnectionEntryEvent::NewVersionAvailable(
226                                    version.into(),
227                                ));
228                            })
229                            .ok();
230                        this.entries.remove(&key);
231                        cx.notify();
232                    })
233                    .ok();
234                    break;
235                }
236            }
237        })
238        .detach();
239
240        cx.spawn({
241            let entry = entry.downgrade();
242            async move |this, cx| {
243                while let Ok(status) = loading_status_rx.recv().await {
244                    let status = status.map(SharedString::from);
245                    let key = key.clone();
246                    let entry = entry.clone();
247                    this.update(cx, move |this, cx| {
248                        if this.entries.get(&key) != entry.upgrade().as_ref() {
249                            return;
250                        }
251
252                        entry
253                            .update(cx, move |_entry, cx| {
254                                cx.emit(AgentConnectionEntryEvent::LoadingStatusChanged(status));
255                            })
256                            .ok();
257                        cx.notify();
258                    })
259                    .ok();
260                }
261            }
262        })
263        .detach();
264
265        entry
266    }
267
268    fn handle_agent_servers_updated(
269        &mut self,
270        store: Entity<AgentServerStore>,
271        _: &AgentServersUpdated,
272        cx: &mut Context<Self>,
273    ) {
274        let store = store.read(cx);
275        self.entries.retain(|key, _| match key {
276            Agent::NativeAgent => true,
277            Agent::Custom { id } => store.external_agents.contains_key(id),
278            #[cfg(any(test, feature = "test-support"))]
279            Agent::Stub => true,
280        });
281        cx.notify();
282    }
283
284    fn start_connection(
285        &self,
286        server: Rc<dyn AgentServer>,
287        cx: &mut Context<Self>,
288    ) -> (
289        Receiver<Option<String>>,
290        Receiver<Option<String>>,
291        Task<Result<AgentConnectedState, LoadError>>,
292    ) {
293        let (new_version_tx, new_version_rx) = watch::channel::<Option<String>>(None);
294        let (loading_status_tx, loading_status_rx) = watch::channel::<Option<String>>(None);
295
296        let agent_server_store = self.project.read(cx).agent_server_store().clone();
297        let delegate = AgentServerDelegate::new(
298            agent_server_store,
299            Some(new_version_tx),
300            Some(loading_status_tx),
301        );
302
303        let connect_task = server.connect(delegate, self.project.clone(), cx);
304        let connect_task = cx.spawn(async move |_this, _cx| match connect_task.await {
305            Ok(connection) => Ok(AgentConnectedState { connection }),
306            Err(err) => match err.downcast::<LoadError>() {
307                Ok(load_error) => Err(load_error),
308                Err(err) => Err(LoadError::Other(SharedString::from(err.to_string()))),
309            },
310        });
311        (new_version_rx, loading_status_rx, connect_task)
312    }
313}
314