tenant.openagents/omega

1use crate::ns_string;
2use anyhow::{Result, anyhow};
3use block::ConcreteBlock;
4use cocoa::{
5    base::{YES, id, nil},
6    foundation::NSArray,
7};
8use collections::HashMap;
9use core_foundation::base::TCFType;
10use core_graphics::display::{
11    CGDirectDisplayID, CGDisplayCopyDisplayMode, CGDisplayModeGetPixelHeight,
12    CGDisplayModeGetPixelWidth, CGDisplayModeRelease,
13};
14use ctor::ctor;
15use futures::channel::oneshot;
16use gpui::{
17    DevicePixels, ForegroundExecutor, ScreenCaptureFrame, ScreenCaptureSource, ScreenCaptureStream,
18    SharedString, SourceMetadata, size,
19};
20use media::core_media::{CMSampleBuffer, CMSampleBufferRef};
21use metal::NSInteger;
22use objc::{
23    class,
24    declare::ClassDecl,
25    msg_send,
26    runtime::{Class, Object, Sel},
27    sel, sel_impl,
28};
29use std::{cell::RefCell, ffi::c_void, mem, ptr, rc::Rc};
30
31use crate::NSStringExt;
32
33#[derive(Clone)]
34pub struct MacScreenCaptureSource {
35    sc_display: id,
36    meta: Option<ScreenMeta>,
37}
38
39pub struct MacScreenCaptureStream {
40    sc_stream: id,
41    sc_stream_output: id,
42    meta: SourceMetadata,
43}
44
45static mut DELEGATE_CLASS: *const Class = ptr::null();
46static mut OUTPUT_CLASS: *const Class = ptr::null();
47const FRAME_CALLBACK_IVAR: &str = "frame_callback";
48
49#[allow(non_upper_case_globals)]
50const SCStreamOutputTypeScreen: NSInteger = 0;
51
52impl ScreenCaptureSource for MacScreenCaptureSource {
53    fn metadata(&self) -> Result<SourceMetadata> {
54        let (display_id, size) = unsafe {
55            let display_id: CGDirectDisplayID = msg_send![self.sc_display, displayID];
56            let display_mode_ref = CGDisplayCopyDisplayMode(display_id);
57            let width = CGDisplayModeGetPixelWidth(display_mode_ref);
58            let height = CGDisplayModeGetPixelHeight(display_mode_ref);
59            CGDisplayModeRelease(display_mode_ref);
60
61            (
62                display_id,
63                size(DevicePixels(width as i32), DevicePixels(height as i32)),
64            )
65        };
66        let (label, is_main) = self
67            .meta
68            .clone()
69            .map(|meta| (meta.label, meta.is_main))
70            .unzip();
71
72        Ok(SourceMetadata {
73            id: display_id as u64,
74            label,
75            is_main,
76            resolution: size,
77        })
78    }
79
80    fn stream(
81        &self,
82        _foreground_executor: &ForegroundExecutor,
83        frame_callback: Box<dyn Fn(ScreenCaptureFrame) + Send>,
84    ) -> oneshot::Receiver<Result<Box<dyn ScreenCaptureStream>>> {
85        unsafe {
86            let stream: id = msg_send![class!(SCStream), alloc];
87            let filter: id = msg_send![class!(SCContentFilter), alloc];
88            let configuration: id = msg_send![class!(SCStreamConfiguration), alloc];
89            let delegate: id = msg_send![DELEGATE_CLASS, alloc];
90            let output: id = msg_send![OUTPUT_CLASS, alloc];
91
92            let excluded_windows = NSArray::array(nil);
93            let filter: id = msg_send![filter, initWithDisplay:self.sc_display excludingWindows:excluded_windows];
94            let configuration: id = msg_send![configuration, init];
95            let _: id = msg_send![configuration, setScalesToFit: true];
96            let _: id = msg_send![configuration, setPixelFormat: 0x42475241];
97            // let _: id = msg_send![configuration, setShowsCursor: false];
98            // let _: id = msg_send![configuration, setCaptureResolution: 3];
99            let delegate: id = msg_send![delegate, init];
100            let output: id = msg_send![output, init];
101
102            output.as_mut().unwrap().set_ivar(
103                FRAME_CALLBACK_IVAR,
104                Box::into_raw(Box::new(frame_callback)) as *mut c_void,
105            );
106
107            let meta = self.metadata().unwrap();
108            let _: id = msg_send![configuration, setWidth: meta.resolution.width.0 as i64];
109            let _: id = msg_send![configuration, setHeight: meta.resolution.height.0 as i64];
110            let stream: id = msg_send![stream, initWithFilter:filter configuration:configuration delegate:delegate];
111
112            // Stream contains filter, configuration, and delegate internally so we release them here
113            // to prevent a memory leak when steam is dropped
114            let _: () = msg_send![filter, release];
115            let _: () = msg_send![configuration, release];
116            let _: () = msg_send![delegate, release];
117
118            let (tx, rx) = oneshot::channel();
119
120            let mut error: id = nil;
121            let _: () = msg_send![stream, addStreamOutput:output type:SCStreamOutputTypeScreen sampleHandlerQueue:0 error:&mut error as *mut id];
122            if error != nil {
123                let message: id = msg_send![error, localizedDescription];
124                let _: () = msg_send![stream, release];
125                let _: () = msg_send![output, release];
126                tx.send(Err(anyhow!("failed to add stream output {message:?}")))
127                    .ok();
128                return rx;
129            }
130
131            let tx = Rc::new(RefCell::new(Some(tx)));
132            let handler = ConcreteBlock::new({
133                move |error: id| {
134                    let result = if error == nil {
135                        let stream = MacScreenCaptureStream {
136                            meta: meta.clone(),
137                            sc_stream: stream,
138                            sc_stream_output: output,
139                        };
140                        Ok(Box::new(stream) as Box<dyn ScreenCaptureStream>)
141                    } else {
142                        let _: () = msg_send![stream, release];
143                        let _: () = msg_send![output, release];
144                        let message: id = msg_send![error, localizedDescription];
145                        Err(anyhow!("failed to start screen capture stream {message:?}"))
146                    };
147                    if let Some(tx) = tx.borrow_mut().take() {
148                        tx.send(result).ok();
149                    }
150                }
151            });
152            let handler = handler.copy();
153            let _: () = msg_send![stream, startCaptureWithCompletionHandler:handler];
154            rx
155        }
156    }
157}
158
159impl Drop for MacScreenCaptureSource {
160    fn drop(&mut self) {
161        unsafe {
162            let _: () = msg_send![self.sc_display, release];
163        }
164    }
165}
166
167impl ScreenCaptureStream for MacScreenCaptureStream {
168    fn metadata(&self) -> Result<SourceMetadata> {
169        Ok(self.meta.clone())
170    }
171}
172
173impl Drop for MacScreenCaptureStream {
174    fn drop(&mut self) {
175        unsafe {
176            let mut error: id = nil;
177            let _: () = msg_send![self.sc_stream, removeStreamOutput:self.sc_stream_output type:SCStreamOutputTypeScreen error:&mut error as *mut _];
178            if error != nil {
179                let message: id = msg_send![error, localizedDescription];
180                log::error!("failed to add stream  output {message:?}");
181            }
182
183            let handler = ConcreteBlock::new(move |error: id| {
184                if error != nil {
185                    let message: id = msg_send![error, localizedDescription];
186                    log::error!("failed to stop screen capture stream {message:?}");
187                }
188            });
189            let block = handler.copy();
190            let _: () = msg_send![self.sc_stream, stopCaptureWithCompletionHandler:block];
191            let _: () = msg_send![self.sc_stream, release];
192            let _: () = msg_send![self.sc_stream_output, release];
193        }
194    }
195}
196
197#[derive(Clone)]
198struct ScreenMeta {
199    label: SharedString,
200    // Is this the screen with menu bar?
201    is_main: bool,
202}
203
204unsafe fn screen_id_to_human_label() -> HashMap<CGDirectDisplayID, ScreenMeta> {
205    let screens: id = msg_send![class!(NSScreen), screens];
206    let count: usize = msg_send![screens, count];
207    let mut map = HashMap::default();
208    let screen_number_key = unsafe { ns_string("NSScreenNumber") };
209    for i in 0..count {
210        let screen: id = msg_send![screens, objectAtIndex: i];
211        let device_desc: id = msg_send![screen, deviceDescription];
212        if device_desc == nil {
213            continue;
214        }
215
216        let nsnumber: id = msg_send![device_desc, objectForKey: screen_number_key];
217        if nsnumber == nil {
218            continue;
219        }
220
221        let screen_id: u32 = msg_send![nsnumber, unsignedIntValue];
222
223        let name: id = msg_send![screen, localizedName];
224        if name != nil {
225            let cstr: *const std::os::raw::c_char = msg_send![name, UTF8String];
226            let rust_str = unsafe {
227                std::ffi::CStr::from_ptr(cstr)
228                    .to_string_lossy()
229                    .into_owned()
230            };
231            map.insert(
232                screen_id,
233                ScreenMeta {
234                    label: rust_str.into(),
235                    is_main: i == 0,
236                },
237            );
238        }
239    }
240    map
241}
242
243pub(crate) fn get_sources() -> oneshot::Receiver<Result<Vec<Rc<dyn ScreenCaptureSource>>>> {
244    unsafe {
245        let (tx, rx) = oneshot::channel();
246        let tx = Rc::new(RefCell::new(Some(tx)));
247        let screen_id_to_label = screen_id_to_human_label();
248        let block = ConcreteBlock::new(move |shareable_content: id, error: id| {
249            let Some(tx) = tx.borrow_mut().take() else {
250                return;
251            };
252
253            let result = if error == nil {
254                let displays: id = msg_send![shareable_content, displays];
255                let mut result = Vec::new();
256                for i in 0..displays.count() {
257                    let display = displays.objectAtIndex(i);
258                    let id: CGDirectDisplayID = msg_send![display, displayID];
259                    let meta = screen_id_to_label.get(&id).cloned();
260                    let source = MacScreenCaptureSource {
261                        sc_display: msg_send![display, retain],
262                        meta,
263                    };
264                    result.push(Rc::new(source) as Rc<dyn ScreenCaptureSource>);
265                }
266                Ok(result)
267            } else {
268                let msg: id = msg_send![error, localizedDescription];
269                Err(anyhow!(
270                    "Screen share failed: {:?}",
271                    NSStringExt::to_str(&msg)
272                ))
273            };
274            tx.send(result).ok();
275        });
276        let block = block.copy();
277
278        let _: () = msg_send![
279            class!(SCShareableContent),
280            getShareableContentExcludingDesktopWindows:YES
281                                   onScreenWindowsOnly:YES
282                                     completionHandler:block];
283        rx
284    }
285}
286
287#[ctor(unsafe)]
288unsafe fn build_classes() {
289    let mut decl = ClassDecl::new("GPUIStreamDelegate", class!(NSObject)).unwrap();
290    unsafe {
291        decl.add_method(
292            sel!(outputVideoEffectDidStartForStream:),
293            output_video_effect_did_start_for_stream as extern "C" fn(&Object, Sel, id),
294        );
295        decl.add_method(
296            sel!(outputVideoEffectDidStopForStream:),
297            output_video_effect_did_stop_for_stream as extern "C" fn(&Object, Sel, id),
298        );
299        decl.add_method(
300            sel!(stream:didStopWithError:),
301            stream_did_stop_with_error as extern "C" fn(&Object, Sel, id, id),
302        );
303        DELEGATE_CLASS = decl.register();
304
305        let mut decl = ClassDecl::new("GPUIStreamOutput", class!(NSObject)).unwrap();
306        decl.add_method(
307            sel!(stream:didOutputSampleBuffer:ofType:),
308            stream_did_output_sample_buffer_of_type
309                as extern "C" fn(&Object, Sel, id, id, NSInteger),
310        );
311        decl.add_ivar::<*mut c_void>(FRAME_CALLBACK_IVAR);
312
313        OUTPUT_CLASS = decl.register();
314    }
315}
316
317extern "C" fn output_video_effect_did_start_for_stream(_this: &Object, _: Sel, _stream: id) {}
318
319extern "C" fn output_video_effect_did_stop_for_stream(_this: &Object, _: Sel, _stream: id) {}
320
321extern "C" fn stream_did_stop_with_error(_this: &Object, _: Sel, _stream: id, _error: id) {}
322
323extern "C" fn stream_did_output_sample_buffer_of_type(
324    this: &Object,
325    _: Sel,
326    _stream: id,
327    sample_buffer: id,
328    buffer_type: NSInteger,
329) {
330    if buffer_type != SCStreamOutputTypeScreen {
331        return;
332    }
333
334    unsafe {
335        let sample_buffer = sample_buffer as CMSampleBufferRef;
336        let sample_buffer = CMSampleBuffer::wrap_under_get_rule(sample_buffer);
337        if let Some(buffer) = sample_buffer.image_buffer() {
338            let callback: Box<Box<dyn Fn(ScreenCaptureFrame)>> =
339                Box::from_raw(*this.get_ivar::<*mut c_void>(FRAME_CALLBACK_IVAR) as *mut _);
340            callback(ScreenCaptureFrame(buffer));
341            mem::forget(callback);
342        }
343    }
344}
345