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}