tenant.openagents/omega

1use std::sync::Arc;
2
3use futures::{FutureExt as _, channel::oneshot, future::Shared};
4use gpui::{App, AppContext as _, AsyncApp, Global, Task};
5use omega_identity::{CustodyError, CustodyState, IdentityInspection, IdentityService};
6use workspace::AppState;
7
8use crate::show_onboarding_view;
9
10type StartupInspection = Result<IdentityInspection, Arc<CustodyError>>;
11type StartupCompletion = Result<(), Arc<anyhow::Error>>;
12
13trait IdentityStartupBackend: Send + Sync {
14    fn inspect_for_process_start(&self) -> Result<IdentityInspection, CustodyError>;
15}
16
17struct SystemIdentityStartupBackend {
18    service: IdentityService,
19}
20
21impl SystemIdentityStartupBackend {
22    fn new() -> Self {
23        Self {
24            service: IdentityService::system(*app_identity::CHANNEL),
25        }
26    }
27}
28
29impl IdentityStartupBackend for SystemIdentityStartupBackend {
30    fn inspect_for_process_start(&self) -> Result<IdentityInspection, CustodyError> {
31        self.service.inspect_for_process_start()
32    }
33}
34
35struct IdentityStartupCoordinator {
36    inspection: Shared<Task<StartupInspection>>,
37    completion: Shared<Task<StartupCompletion>>,
38    completion_sender: Option<oneshot::Sender<StartupCompletion>>,
39    onboarding_open: bool,
40    terminal: Option<StartupCompletion>,
41}
42
43impl Global for IdentityStartupCoordinator {}
44
45impl IdentityStartupCoordinator {
46    fn install(cx: &mut App) {
47        if cx.has_global::<Self>() {
48            return;
49        }
50        Self::install_with_backend(Arc::new(SystemIdentityStartupBackend::new()), cx);
51    }
52
53    fn install_with_backend(backend: Arc<dyn IdentityStartupBackend>, cx: &mut App) {
54        if cx.has_global::<Self>() {
55            return;
56        }
57
58        let inspection = cx
59            .background_spawn(async move { backend.inspect_for_process_start().map_err(Arc::new) })
60            .shared();
61        let (completion_sender, completion_receiver) = oneshot::channel();
62        let completion = cx
63            .spawn(async move |_| match completion_receiver.await {
64                Ok(completion) => completion,
65                Err(_) => Err(Arc::new(anyhow::anyhow!(
66                    "identity startup completion channel closed before release"
67                ))),
68            })
69            .shared();
70
71        cx.set_global(Self {
72            inspection,
73            completion,
74            completion_sender: Some(completion_sender),
75            onboarding_open: false,
76            terminal: None,
77        });
78    }
79
80    fn claim_onboarding(&mut self) -> bool {
81        if self.terminal.is_some() || self.onboarding_open {
82            return false;
83        }
84        self.onboarding_open = true;
85        true
86    }
87
88    fn onboarding_opened(cx: &mut App) {
89        if cx.has_global::<Self>() {
90            cx.global_mut::<Self>().onboarding_open = true;
91        }
92    }
93
94    fn onboarding_closed(cx: &mut App) {
95        if cx.has_global::<Self>() {
96            cx.global_mut::<Self>().onboarding_open = false;
97        }
98    }
99
100    fn finish(completion: StartupCompletion, cx: &mut App) {
101        if !cx.has_global::<Self>() {
102            return;
103        }
104        let coordinator = cx.global_mut::<Self>();
105        if coordinator.terminal.is_some() {
106            return;
107        }
108
109        coordinator.terminal = Some(completion.clone());
110        if let Some(sender) = coordinator.completion_sender.take() {
111            if sender.send(completion).is_err() {
112                zlog::error!("identity startup waiters disappeared before release");
113            }
114        }
115    }
116}
117
118pub async fn await_identity_ready(
119    app_state: Arc<AppState>,
120    cx: &mut AsyncApp,
121) -> anyhow::Result<()> {
122    let (inspection, completion, terminal) = cx.update(|cx| {
123        IdentityStartupCoordinator::install(cx);
124        let coordinator = cx.global::<IdentityStartupCoordinator>();
125        (
126            coordinator.inspection.clone(),
127            coordinator.completion.clone(),
128            coordinator.terminal.is_some(),
129        )
130    });
131    if terminal {
132        return completion
133            .await
134            .map_err(|error| anyhow::anyhow!("{error:#}"));
135    };
136
137    let needs_onboarding = match inspection.await {
138        Ok(inspection) => inspection.custody.state != CustodyState::Ready,
139        Err(error) => {
140            zlog::error!("identity startup inspection failed: {error}");
141            true
142        }
143    };
144
145    if !needs_onboarding {
146        cx.update(|cx| IdentityStartupCoordinator::finish(Ok(()), cx));
147        return Ok(());
148    }
149
150    let should_open = cx.update(|cx| {
151        let coordinator = cx.global_mut::<IdentityStartupCoordinator>();
152        coordinator.claim_onboarding()
153    });
154    if should_open {
155        let open_task = cx.update(|cx| show_onboarding_view(app_state, cx));
156        cx.spawn(async move |cx| {
157            if let Err(error) = open_task.await {
158                zlog::error!("failed to open identity onboarding: {error:#}");
159                cx.update(|cx| {
160                    IdentityStartupCoordinator::onboarding_closed(cx);
161                    IdentityStartupCoordinator::finish(Err(Arc::new(error)), cx);
162                });
163            }
164        })
165        .detach();
166    }
167
168    completion
169        .await
170        .map_err(|error| anyhow::anyhow!("{error:#}"))
171}
172
173pub(crate) fn onboarding_opened(cx: &mut App) {
174    IdentityStartupCoordinator::onboarding_opened(cx);
175}
176
177pub(crate) fn onboarding_closed(cx: &mut App) {
178    IdentityStartupCoordinator::onboarding_closed(cx);
179}
180
181pub(crate) fn release_identity_waiters(cx: &mut App) {
182    IdentityStartupCoordinator::finish(Ok(()), cx);
183}
184
185#[cfg(test)]
186mod tests {
187    use std::sync::atomic::{AtomicUsize, Ordering};
188
189    use gpui::TestAppContext;
190    use omega_identity::{CustodyResult, RecoveryProtectionState, RecoveryProtectionStatus};
191
192    use super::*;
193
194    struct FakeBackend {
195        calls: Arc<AtomicUsize>,
196        state: CustodyState,
197    }
198
199    impl IdentityStartupBackend for FakeBackend {
200        fn inspect_for_process_start(&self) -> Result<IdentityInspection, CustodyError> {
201            self.calls.fetch_add(1, Ordering::SeqCst);
202            Ok(IdentityInspection {
203                custody: CustodyResult {
204                    state: self.state,
205                    identity: None,
206                    receipt_ref: None,
207                },
208                pending_transaction: None,
209                conflict: None,
210                recovery_protection: RecoveryProtectionStatus {
211                    state: RecoveryProtectionState::NotApplicable,
212                    record: None,
213                },
214            })
215        }
216    }
217
218    #[gpui::test]
219    fn startup_inspection_is_shared_by_concurrent_callers(cx: &mut TestAppContext) {
220        let calls = Arc::new(AtomicUsize::new(0));
221        let backend = Arc::new(FakeBackend {
222            calls: calls.clone(),
223            state: CustodyState::Ready,
224        });
225        let completed = Arc::new(AtomicUsize::new(0));
226
227        cx.update(|cx| {
228            IdentityStartupCoordinator::install_with_backend(backend, cx);
229            let inspection = cx.global::<IdentityStartupCoordinator>().inspection.clone();
230            for _ in 0..2 {
231                let inspection = inspection.clone();
232                let completed = completed.clone();
233                cx.spawn(async move |_| {
234                    inspection.await.expect("startup inspection");
235                    completed.fetch_add(1, Ordering::SeqCst);
236                })
237                .detach();
238            }
239        });
240        cx.run_until_parked();
241
242        assert_eq!(calls.load(Ordering::SeqCst), 1);
243        assert_eq!(completed.load(Ordering::SeqCst), 2);
244    }
245
246    #[gpui::test]
247    fn release_wakes_every_waiter_and_is_idempotent(cx: &mut TestAppContext) {
248        let calls = Arc::new(AtomicUsize::new(0));
249        let backend = Arc::new(FakeBackend {
250            calls,
251            state: CustodyState::Absent,
252        });
253        let completed = Arc::new(AtomicUsize::new(0));
254
255        cx.update(|cx| {
256            IdentityStartupCoordinator::install_with_backend(backend, cx);
257            let completion = cx.global::<IdentityStartupCoordinator>().completion.clone();
258            for _ in 0..3 {
259                let completion = completion.clone();
260                let completed = completed.clone();
261                cx.spawn(async move |_| {
262                    completion.await.expect("startup release");
263                    completed.fetch_add(1, Ordering::SeqCst);
264                })
265                .detach();
266            }
267            IdentityStartupCoordinator::finish(Ok(()), cx);
268            IdentityStartupCoordinator::finish(Ok(()), cx);
269        });
270        cx.run_until_parked();
271
272        assert_eq!(completed.load(Ordering::SeqCst), 3);
273        cx.update(|cx| {
274            assert_eq!(
275                cx.global::<IdentityStartupCoordinator>()
276                    .terminal
277                    .as_ref()
278                    .map(Result::is_ok),
279                Some(true)
280            );
281        });
282    }
283
284    #[gpui::test]
285    fn failure_wakes_current_and_future_waiters(cx: &mut TestAppContext) {
286        let calls = Arc::new(AtomicUsize::new(0));
287        let backend = Arc::new(FakeBackend {
288            calls,
289            state: CustodyState::Absent,
290        });
291        let failures = Arc::new(AtomicUsize::new(0));
292
293        cx.update(|cx| {
294            IdentityStartupCoordinator::install_with_backend(backend, cx);
295            let completion = cx.global::<IdentityStartupCoordinator>().completion.clone();
296            for _ in 0..2 {
297                let completion = completion.clone();
298                let failures = failures.clone();
299                cx.spawn(async move |_| {
300                    if completion.await.is_err() {
301                        failures.fetch_add(1, Ordering::SeqCst);
302                    }
303                })
304                .detach();
305            }
306            IdentityStartupCoordinator::finish(
307                Err(Arc::new(anyhow::anyhow!("failed to open onboarding"))),
308                cx,
309            );
310        });
311        cx.run_until_parked();
312
313        cx.update(|cx| {
314            let completion = cx.global::<IdentityStartupCoordinator>().completion.clone();
315            let failures = failures.clone();
316            cx.spawn(async move |_| {
317                if completion.await.is_err() {
318                    failures.fetch_add(1, Ordering::SeqCst);
319                }
320            })
321            .detach();
322        });
323        cx.run_until_parked();
324
325        assert_eq!(failures.load(Ordering::SeqCst), 3);
326        cx.update(|cx| {
327            assert_eq!(
328                cx.global::<IdentityStartupCoordinator>()
329                    .terminal
330                    .as_ref()
331                    .map(Result::is_err),
332                Some(true)
333            );
334        });
335    }
336
337    #[gpui::test]
338    fn closing_onboarding_allows_one_reopen_without_releasing(cx: &mut TestAppContext) {
339        let calls = Arc::new(AtomicUsize::new(0));
340        let backend = Arc::new(FakeBackend {
341            calls,
342            state: CustodyState::Absent,
343        });
344
345        cx.update(|cx| {
346            IdentityStartupCoordinator::install_with_backend(backend, cx);
347            let coordinator = cx.global_mut::<IdentityStartupCoordinator>();
348            assert!(coordinator.claim_onboarding());
349            assert!(!coordinator.claim_onboarding());
350            IdentityStartupCoordinator::onboarding_closed(cx);
351            assert!(
352                cx.global_mut::<IdentityStartupCoordinator>()
353                    .claim_onboarding()
354            );
355            assert!(cx.global::<IdentityStartupCoordinator>().terminal.is_none());
356        });
357    }
358}
359