tenant.openagents/omega

1use anyhow::Result;
2use futures::StreamExt as _;
3use futures::channel::oneshot;
4use gpui::{AsyncApp, ScreenCaptureStream};
5use livekit::track;
6use livekit::webrtc::{
7    prelude::NV12Buffer,
8    video_frame::{VideoFrame, VideoRotation},
9    video_source::{RtcVideoSource, VideoResolution, native::NativeVideoSource},
10};
11use std::sync::{
12    Arc,
13    atomic::{AtomicBool, AtomicU64, Ordering},
14};
15
16static NEXT_WAYLAND_SHARE_ID: AtomicU64 = AtomicU64::new(1);
17const PIPEWIRE_TIMEOUT_S: u64 = 30;
18
19pub struct WaylandScreenCaptureStream {
20    id: u64,
21    stop_flag: Arc<AtomicBool>,
22    _capture_task: gpui::Task<()>,
23}
24
25impl WaylandScreenCaptureStream {
26    pub fn new(stop_flag: Arc<AtomicBool>, capture_task: gpui::Task<()>) -> Self {
27        Self {
28            id: NEXT_WAYLAND_SHARE_ID.fetch_add(1, Ordering::Relaxed),
29            stop_flag,
30            _capture_task: capture_task,
31        }
32    }
33}
34
35impl ScreenCaptureStream for WaylandScreenCaptureStream {
36    fn metadata(&self) -> Result<gpui::SourceMetadata> {
37        Ok(gpui::SourceMetadata {
38            id: self.id,
39            label: None,
40            is_main: None,
41            resolution: gpui::size(gpui::DevicePixels(1), gpui::DevicePixels(1)),
42        })
43    }
44}
45
46impl Drop for WaylandScreenCaptureStream {
47    fn drop(&mut self) {
48        self.stop_flag.store(true, Ordering::Release);
49    }
50}
51
52pub(crate) async fn start_wayland_desktop_capture(
53    cx: &mut AsyncApp,
54) -> Result<(
55    crate::LocalVideoTrack,
56    Arc<AtomicBool>,
57    gpui::Task<()>,
58    oneshot::Receiver<()>,
59)> {
60    use futures::channel::mpsc;
61    use gpui::FutureExt as _;
62    use libwebrtc::desktop_capturer::{
63        CaptureError, DesktopCaptureSourceType, DesktopCapturer, DesktopCapturerOptions,
64        DesktopFrame,
65    };
66    use libwebrtc::native::yuv_helper::argb_to_nv12;
67    use std::time::Duration;
68    use webrtc_sys::webrtc::ffi as webrtc_ffi;
69
70    fn webrtc_log_callback(message: String, severity: webrtc_ffi::LoggingSeverity) {
71        match severity {
72            webrtc_ffi::LoggingSeverity::Error => log::error!("[webrtc] {}", message.trim()),
73            _ => log::debug!("[webrtc] {}", message.trim()),
74        }
75    }
76
77    let _webrtc_log_sink = webrtc_ffi::new_log_sink(webrtc_log_callback);
78    log::debug!("Wayland desktop capture: WebRTC internal logging enabled");
79
80    let stop_flag = Arc::new(AtomicBool::new(false));
81    let (mut video_source_tx, mut video_source_rx) = mpsc::channel::<NativeVideoSource>(1);
82    let (failure_tx, failure_rx) = oneshot::channel::<()>();
83
84    let mut options = DesktopCapturerOptions::new(DesktopCaptureSourceType::Generic);
85    options.set_include_cursor(true);
86    let mut capturer = DesktopCapturer::new(options).ok_or_else(|| {
87        anyhow::anyhow!(
88            "Failed to create desktop capturer. \
89             Check that xdg-desktop-portal is installed and running."
90        )
91    })?;
92
93    let permanent_error = Arc::new(AtomicBool::new(false));
94    let stop_cb = stop_flag.clone();
95    let permanent_error_cb = permanent_error.clone();
96    capturer.start_capture(None, {
97        let mut video_source: Option<NativeVideoSource> = None;
98        let mut current_width: u32 = 0;
99        let mut current_height: u32 = 0;
100        let mut video_frame = VideoFrame {
101            rotation: VideoRotation::VideoRotation0,
102            buffer: NV12Buffer::new(1, 1),
103            timestamp_us: 0,
104        };
105
106        move |result: Result<DesktopFrame, CaptureError>| {
107            let frame = match result {
108                Ok(frame) => frame,
109                Err(CaptureError::Temporary) => return,
110                Err(CaptureError::Permanent) => {
111                    log::error!("Wayland desktop capture encountered a permanent error");
112                    permanent_error_cb.store(true, Ordering::Release);
113                    stop_cb.store(true, Ordering::Release);
114                    return;
115                }
116            };
117
118            let width = frame.width() as u32;
119            let height = frame.height() as u32;
120            if width != current_width || height != current_height {
121                current_width = width;
122                current_height = height;
123                video_frame.buffer = NV12Buffer::new(width, height);
124            }
125
126            let (stride_y, stride_uv) = video_frame.buffer.strides();
127            let (data_y, data_uv) = video_frame.buffer.data_mut();
128            argb_to_nv12(
129                frame.data(),
130                frame.stride(),
131                data_y,
132                stride_y,
133                data_uv,
134                stride_uv,
135                width as i32,
136                height as i32,
137            );
138
139            if let Some(source) = &video_source {
140                source.capture_frame(&video_frame);
141            } else {
142                let source = NativeVideoSource::new(VideoResolution { width, height }, true);
143                source.capture_frame(&video_frame);
144                video_source_tx.try_send(source.clone()).ok();
145                video_source = Some(source);
146            }
147        }
148    });
149
150    log::info!("Wayland desktop capture: starting capture loop");
151
152    let stop = stop_flag.clone();
153    let tokio_task = gpui_tokio::Tokio::spawn(cx, async move {
154        loop {
155            if stop.load(Ordering::Acquire) {
156                break;
157            }
158            capturer.capture_frame();
159            tokio::time::sleep(Duration::from_millis(33)).await;
160        }
161        drop(capturer);
162
163        if permanent_error.load(Ordering::Acquire) {
164            log::error!("Wayland screen capture ended due to a permanent capture error");
165            let _ = failure_tx.send(());
166        }
167    });
168
169    let capture_task = cx.background_executor().spawn(async move {
170        if let Err(error) = tokio_task.await {
171            log::error!("Wayland capture task failed: {error}");
172        }
173    });
174
175    let executor = cx.background_executor().clone();
176    let video_source = video_source_rx
177        .next()
178        .with_timeout(Duration::from_secs(PIPEWIRE_TIMEOUT_S), &executor)
179        .await
180        .map_err(|_| {
181            stop_flag.store(true, Ordering::Relaxed);
182            log::error!("Wayland desktop capture timed out.");
183            anyhow::anyhow!(
184                "Screen sharing timed out waiting for the first frame. \
185                 Check that xdg-desktop-portal and PipeWire are running, \
186                 and that your portal backend matches your compositor."
187            )
188        })?
189        .ok_or_else(|| {
190            stop_flag.store(true, Ordering::Relaxed);
191            anyhow::anyhow!(
192                "Screen sharing was canceled or the portal denied permission. \
193                 You can try again from the screen share button."
194            )
195        })?;
196
197    let track = super::LocalVideoTrack(track::LocalVideoTrack::create_video_track(
198        "screen share",
199        RtcVideoSource::Native(video_source),
200    ));
201
202    Ok((track, stop_flag, capture_task, failure_rx))
203}
204