tenant.openagents/omega

1use gpui::{
2    PlatformDispatcher, Priority, PriorityQueueReceiver, PriorityQueueSender, RunnableVariant,
3};
4use std::sync::Arc;
5use std::sync::atomic::AtomicI32;
6use std::time::Duration;
7use wasm_bindgen::prelude::*;
8use web_time::Instant;
9
10#[cfg(feature = "multithreaded")]
11const MIN_BACKGROUND_THREADS: usize = 2;
12
13#[cfg(feature = "multithreaded")]
14fn shared_memory_supported() -> bool {
15    let global = js_sys::global();
16    let has_shared_array_buffer =
17        js_sys::Reflect::has(&global, &JsValue::from_str("SharedArrayBuffer")).unwrap_or(false);
18    let has_atomics = js_sys::Reflect::has(&global, &JsValue::from_str("Atomics")).unwrap_or(false);
19    let memory = js_sys::WebAssembly::Memory::from(wasm_bindgen::memory());
20    let buffer = memory.buffer();
21    let is_shared_buffer = buffer.is_instance_of::<js_sys::SharedArrayBuffer>();
22    has_shared_array_buffer && has_atomics && is_shared_buffer
23}
24
25enum MainThreadItem {
26    Runnable(RunnableVariant),
27    Delayed {
28        runnable: RunnableVariant,
29        millis: i32,
30    },
31    // TODO-Wasm: Shouldn't these run on their own dedicated thread?
32    RealtimeFunction(Box<dyn FnOnce() + Send>),
33}
34
35struct MainThreadMailbox {
36    sender: PriorityQueueSender<MainThreadItem>,
37    receiver: parking_lot::Mutex<PriorityQueueReceiver<MainThreadItem>>,
38    signal: AtomicI32,
39}
40
41impl MainThreadMailbox {
42    fn new() -> Self {
43        let (sender, receiver) = PriorityQueueReceiver::new();
44        Self {
45            sender,
46            receiver: parking_lot::Mutex::new(receiver),
47            signal: AtomicI32::new(0),
48        }
49    }
50
51    fn post(&self, priority: Priority, item: MainThreadItem) {
52        if self.sender.spin_send(priority, item).is_err() {
53            log::error!("MainThreadMailbox::send failed: receiver disconnected");
54        }
55
56        // TODO-Wasm: Verify this lock-free protocol
57        let view = self.signal_view();
58        js_sys::Atomics::store(&view, 0, 1).ok();
59        js_sys::Atomics::notify(&view, 0).ok();
60    }
61
62    fn drain(&self, window: &web_sys::Window) {
63        let mut receiver = self.receiver.lock();
64        loop {
65            // We need these `spin` variants because we can't acquire a lock on the main thread.
66            // TODO-WASM: Should we do something different?
67            match receiver.spin_try_pop() {
68                Ok(Some(item)) => execute_on_main_thread(window, item),
69                Ok(None) => break,
70                Err(_) => break,
71            }
72        }
73    }
74
75    fn signal_view(&self) -> js_sys::Int32Array {
76        let byte_offset = self.signal.as_ptr() as u32;
77        let memory = js_sys::WebAssembly::Memory::from(wasm_bindgen::memory());
78        js_sys::Int32Array::new_with_byte_offset_and_length(&memory.buffer(), byte_offset, 1)
79    }
80
81    fn run_waker_loop(self: &Arc<Self>, window: web_sys::Window) {
82        if !shared_memory_supported() {
83            log::warn!("SharedArrayBuffer not available; main thread mailbox waker loop disabled");
84            return;
85        }
86
87        let mailbox = Arc::clone(self);
88        wasm_bindgen_futures::spawn_local(async move {
89            let view = mailbox.signal_view();
90            loop {
91                js_sys::Atomics::store(&view, 0, 0).expect("Atomics.store failed");
92
93                let result = match js_sys::Atomics::wait_async(&view, 0, 0) {
94                    Ok(result) => result,
95                    Err(error) => {
96                        log::error!("Atomics.waitAsync failed: {error:?}");
97                        break;
98                    }
99                };
100
101                let is_async = js_sys::Reflect::get(&result, &JsValue::from_str("async"))
102                    .ok()
103                    .and_then(|v| v.as_bool())
104                    .unwrap_or(false);
105
106                if !is_async {
107                    log::error!("Atomics.waitAsync returned synchronously; waker loop exiting");
108                    break;
109                }
110
111                let promise: js_sys::Promise =
112                    js_sys::Reflect::get(&result, &JsValue::from_str("value"))
113                        .expect("waitAsync result missing 'value'")
114                        .unchecked_into();
115
116                let _ = wasm_bindgen_futures::JsFuture::from(promise).await;
117
118                mailbox.drain(&window);
119            }
120        });
121    }
122}
123
124pub struct WebDispatcher {
125    main_thread_id: std::thread::ThreadId,
126    browser_window: web_sys::Window,
127    background_sender: PriorityQueueSender<RunnableVariant>,
128    main_thread_mailbox: Arc<MainThreadMailbox>,
129    supports_threads: bool,
130    #[cfg(feature = "multithreaded")]
131    _background_threads: Vec<wasm_thread::JoinHandle<()>>,
132}
133
134// Safety: `web_sys::Window` is only accessed from the main thread
135// All other fields are `Send + Sync` by construction.
136unsafe impl Send for WebDispatcher {}
137unsafe impl Sync for WebDispatcher {}
138
139impl WebDispatcher {
140    pub fn new(browser_window: web_sys::Window, allow_threads: bool) -> Self {
141        #[cfg(feature = "multithreaded")]
142        let (background_sender, background_receiver) = PriorityQueueReceiver::new();
143        #[cfg(not(feature = "multithreaded"))]
144        let (background_sender, _) = PriorityQueueReceiver::new();
145
146        let main_thread_mailbox = Arc::new(MainThreadMailbox::new());
147
148        #[cfg(feature = "multithreaded")]
149        let supports_threads = allow_threads && shared_memory_supported();
150        #[cfg(not(feature = "multithreaded"))]
151        let supports_threads = false;
152
153        if supports_threads {
154            main_thread_mailbox.run_waker_loop(browser_window.clone());
155        } else {
156            log::warn!(
157                "SharedArrayBuffer not available; falling back to single-threaded dispatcher"
158            );
159        }
160
161        #[cfg(feature = "multithreaded")]
162        let background_threads = if supports_threads {
163            let thread_count = browser_window
164                .navigator()
165                .hardware_concurrency()
166                .max(MIN_BACKGROUND_THREADS as f64) as usize;
167
168            // TODO-Wasm: Is it bad to have web workers blocking for a long time like this?
169            (0..thread_count)
170                .map(|i| {
171                    let mut receiver = background_receiver.clone();
172                    wasm_thread::Builder::new()
173                        .name(format!("background-worker-{i}"))
174                        .spawn(move || {
175                            loop {
176                                let runnable: RunnableVariant = match receiver.pop() {
177                                    Ok(runnable) => runnable,
178                                    Err(_) => {
179                                        log::info!(
180                                            "background-worker-{i}: channel disconnected, exiting"
181                                        );
182                                        break;
183                                    }
184                                };
185
186                                runnable.run();
187                            }
188                        })
189                        .expect("failed to spawn background worker thread")
190                })
191                .collect::<Vec<_>>()
192        } else {
193            Vec::new()
194        };
195
196        Self {
197            main_thread_id: std::thread::current().id(),
198            browser_window,
199            background_sender,
200            main_thread_mailbox,
201            supports_threads,
202            #[cfg(feature = "multithreaded")]
203            _background_threads: background_threads,
204        }
205    }
206
207    fn on_main_thread(&self) -> bool {
208        std::thread::current().id() == self.main_thread_id
209    }
210}
211
212impl PlatformDispatcher for WebDispatcher {
213    fn is_main_thread(&self) -> bool {
214        self.on_main_thread()
215    }
216
217    fn dispatch(&self, runnable: RunnableVariant, priority: Priority) {
218        if !self.supports_threads {
219            self.dispatch_on_main_thread(runnable, priority);
220            return;
221        }
222
223        let result = if self.on_main_thread() {
224            self.background_sender.spin_send(priority, runnable)
225        } else {
226            self.background_sender.send(priority, runnable)
227        };
228
229        if let Err(error) = result {
230            log::error!("dispatch: failed to send to background queue: {error:?}");
231        }
232    }
233
234    fn dispatch_on_main_thread(&self, runnable: RunnableVariant, priority: Priority) {
235        if self.on_main_thread() {
236            schedule_runnable(&self.browser_window, runnable, priority);
237        } else {
238            self.main_thread_mailbox
239                .post(priority, MainThreadItem::Runnable(runnable));
240        }
241    }
242
243    fn dispatch_after(&self, duration: Duration, runnable: RunnableVariant) {
244        let millis = duration.as_millis().min(i32::MAX as u128) as i32;
245        if self.on_main_thread() {
246            let callback = Closure::once_into_js(move || {
247                runnable.run();
248            });
249            self.browser_window
250                .set_timeout_with_callback_and_timeout_and_arguments_0(
251                    callback.unchecked_ref(),
252                    millis,
253                )
254                .ok();
255        } else {
256            self.main_thread_mailbox
257                .post(Priority::High, MainThreadItem::Delayed { runnable, millis });
258        }
259    }
260
261    fn spawn_realtime(&self, function: Box<dyn FnOnce() + Send>) {
262        if self.on_main_thread() {
263            let callback = Closure::once_into_js(move || {
264                function();
265            });
266            self.browser_window
267                .queue_microtask(callback.unchecked_ref());
268        } else {
269            self.main_thread_mailbox
270                .post(Priority::High, MainThreadItem::RealtimeFunction(function));
271        }
272    }
273
274    fn now(&self) -> Instant {
275        Instant::now()
276    }
277}
278
279fn execute_on_main_thread(window: &web_sys::Window, item: MainThreadItem) {
280    match item {
281        MainThreadItem::Runnable(runnable) => {
282            runnable.run();
283        }
284        MainThreadItem::Delayed { runnable, millis } => {
285            let callback = Closure::once_into_js(move || {
286                runnable.run();
287            });
288            window
289                .set_timeout_with_callback_and_timeout_and_arguments_0(
290                    callback.unchecked_ref(),
291                    millis,
292                )
293                .ok();
294        }
295        MainThreadItem::RealtimeFunction(function) => {
296            function();
297        }
298    }
299}
300
301fn schedule_runnable(window: &web_sys::Window, runnable: RunnableVariant, priority: Priority) {
302    let callback = Closure::once_into_js(move || {
303        runnable.run();
304    });
305    let callback: &js_sys::Function = callback.unchecked_ref();
306
307    match priority {
308        Priority::RealtimeAudio => {
309            window.queue_microtask(callback);
310        }
311        _ => {
312            // TODO-Wasm: this ought to enqueue so we can dequeue with proper priority
313            window
314                .set_timeout_with_callback_and_timeout_and_arguments_0(callback, 0)
315                .ok();
316        }
317    }
318}
319