tenant.openagents/omega

1#![cfg(target_os = "windows")]
2
3use anyhow::{Context as _, Result, bail};
4use gpui::{App, AppContext as _, DismissEvent, Global, actions};
5use std::fmt::Write as _;
6use std::io::{BufRead, BufReader, Write};
7use std::path::{Path, PathBuf};
8use std::time::Duration;
9use util::{ResultExt as _, defer};
10use windows::Win32::Foundation::{VARIANT_BOOL, VARIANT_FALSE};
11use windows::Win32::System::Com::{CLSCTX_INPROC_SERVER, COINIT_MULTITHREADED, CoInitializeEx};
12use windows_core::{BSTR, Interface};
13use workspace::notifications::simple_message_notification::MessageNotification;
14use workspace::notifications::{NotificationId, show_app_notification};
15use wprcontrol::*;
16
17actions!(
18    omega,
19    [
20        /// Starts recording an ETW (Event Tracing for Windows) trace.
21        #[action(deprecated_aliases = ["zed::RecordEtwTrace"])]
22        RecordEtwTrace,
23        /// Starts recording an ETW (Event Tracing for Windows) trace with heap tracing.
24        #[action(deprecated_aliases = ["zed::RecordEtwTraceWithHeapTracing"])]
25        RecordEtwTraceWithHeapTracing,
26        /// Saves an in-progress ETW trace to disk.
27        #[action(deprecated_aliases = ["zed::SaveEtwTrace"])]
28        SaveEtwTrace,
29        /// Cancels an in-progress ETW trace without saving.
30        #[action(deprecated_aliases = ["zed::CancelEtwTrace"])]
31        CancelEtwTrace,
32    ]
33);
34
35struct EtwNotification;
36
37struct EtwSessionHandle {
38    writer: net::OwnedWriteHalf,
39    _listener: net::UnixListener,
40    socket_path: PathBuf,
41}
42
43impl Drop for EtwSessionHandle {
44    fn drop(&mut self) {
45        let _ = std::fs::remove_file(&self.socket_path);
46    }
47}
48
49struct GlobalEtwSession(Option<EtwSessionHandle>);
50
51impl Global for GlobalEtwSession {}
52
53fn has_active_etw_session(cx: &App) -> bool {
54    cx.global::<GlobalEtwSession>().0.is_some()
55}
56
57fn show_etw_notification(cx: &mut App, message: impl Into<gpui::SharedString>) {
58    let message = message.into();
59    show_app_notification(NotificationId::unique::<EtwNotification>(), cx, move |cx| {
60        cx.new(|cx| MessageNotification::new(message.clone(), cx))
61    });
62}
63
64fn show_etw_notification_with_action(
65    cx: &mut App,
66    message: impl Into<gpui::SharedString>,
67    button_label: impl Into<gpui::SharedString>,
68    on_click: impl Fn(&mut gpui::Window, &mut gpui::Context<MessageNotification>)
69    + Send
70    + Sync
71    + 'static,
72) {
73    let message = message.into();
74    let button_label = button_label.into();
75    let on_click = std::sync::Arc::new(on_click);
76    show_app_notification(NotificationId::unique::<EtwNotification>(), cx, move |cx| {
77        let message = message.clone();
78        let button_label = button_label.clone();
79        cx.new(|cx| {
80            MessageNotification::new(message, cx)
81                .primary_message(button_label)
82                .primary_on_click_arc(on_click.clone())
83        })
84    });
85}
86
87fn show_etw_status_notification(cx: &mut App, status: Result<StatusMessage>, output_path: PathBuf) {
88    match status {
89        Ok(StatusMessage::Stopped) => {
90            let display_path = output_path.display().to_string();
91            show_etw_notification_with_action(
92                cx,
93                format!("ETW trace saved to {display_path}"),
94                "Show in File Manager",
95                move |_window, cx| {
96                    cx.reveal_path(&output_path);
97                    cx.emit(DismissEvent);
98                },
99            );
100        }
101        Ok(StatusMessage::TimedOut) => {
102            let display_path = output_path.display().to_string();
103            show_etw_notification_with_action(
104                cx,
105                format!("ETW recording timed out. Trace saved to {display_path}"),
106                "Show in File Manager",
107                move |_window, cx| {
108                    cx.reveal_path(&output_path);
109                    cx.emit(DismissEvent);
110                },
111            );
112        }
113        Ok(StatusMessage::Cancelled) => {
114            show_etw_notification(cx, "ETW recording cancelled");
115        }
116        Ok(_) => {
117            show_etw_notification(cx, "ETW recording ended unexpectedly");
118        }
119        Err(error) => {
120            show_etw_notification(cx, format!("Failed to complete ETW recording: {error:#}"));
121        }
122    }
123}
124
125pub fn init(cx: &mut App) {
126    cx.set_global(GlobalEtwSession(None));
127
128    cx.on_action(|_: &RecordEtwTrace, cx: &mut App| {
129        start_etw_recording(cx, None);
130    });
131
132    cx.on_action(|_: &RecordEtwTraceWithHeapTracing, cx: &mut App| {
133        start_etw_recording(cx, Some(std::process::id()));
134    });
135
136    cx.on_action(|_: &SaveEtwTrace, cx: &mut App| {
137        let session = cx.global_mut::<GlobalEtwSession>().0.as_mut();
138        let Some(session) = session else {
139            show_etw_notification(cx, "No active ETW recording to stop");
140            return;
141        };
142        match send_json(&mut session.writer, &Command::Save) {
143            Ok(()) => {
144                show_etw_notification(cx, "Stopping ETW recording...");
145            }
146            Err(error) => {
147                show_etw_notification(cx, format!("Failed to stop ETW recording: {error:#}"));
148            }
149        }
150    });
151
152    cx.on_action(|_: &CancelEtwTrace, cx: &mut App| {
153        let session = cx.global_mut::<GlobalEtwSession>().0.as_mut();
154        let Some(session) = session else {
155            show_etw_notification(cx, "No active ETW recording to cancel");
156            return;
157        };
158        match send_json(&mut session.writer, &Command::Cancel) {
159            Ok(()) => {
160                show_etw_notification(cx, "Cancelling ETW recording...");
161            }
162            Err(error) => {
163                show_etw_notification(cx, format!("Failed to cancel ETW recording: {error:#}"));
164            }
165        }
166    });
167}
168
169fn start_etw_recording(cx: &mut App, heap_pid: Option<u32>) {
170    if has_active_etw_session(cx) {
171        show_etw_notification(cx, "ETW recording is already in progress");
172        return;
173    }
174    let save_dialog = cx.prompt_for_new_path(&PathBuf::default(), Some("zed-trace.etl"));
175    cx.spawn(async move |cx| {
176        let output_path = match save_dialog.await {
177            Ok(Ok(Some(path))) => path,
178            Ok(Ok(None)) => return,
179            Ok(Err(error)) => {
180                cx.update(|cx| {
181                    show_etw_notification(cx, format!("Failed to pick save location: {error:#}"));
182                });
183                return;
184            }
185            Err(_) => return,
186        };
187
188        let result = cx
189            .background_spawn(async move { launch_etw_recording(heap_pid, &output_path) })
190            .await;
191
192        let EtwSession {
193            output_path,
194            stream,
195            listener,
196            socket_path,
197        } = match result {
198            Ok(session) => session,
199            Err(error) => {
200                cx.update(|cx| {
201                    show_etw_notification(cx, format!("Failed to start ETW recording: {error:#}"));
202                });
203                return;
204            }
205        };
206
207        let (read_half, write_half) = stream.into_inner().into_split();
208
209        cx.spawn(async |cx| {
210            let status = cx
211                .background_spawn(async move {
212                    recv_json(&mut BufReader::new(read_half))
213                        .context("Receive status from subprocess")
214                })
215                .await;
216            cx.update(|cx| {
217                cx.global_mut::<GlobalEtwSession>().0 = None;
218                show_etw_status_notification(cx, status, output_path);
219            });
220        })
221        .detach();
222
223        cx.update(|cx| {
224            cx.global_mut::<GlobalEtwSession>().0 = Some(EtwSessionHandle {
225                writer: write_half,
226                _listener: listener,
227                socket_path,
228            });
229            show_etw_notification(cx, "ETW recording started");
230        });
231    })
232    .detach();
233}
234
235const RECORDING_TIMEOUT: Duration = Duration::from_secs(60);
236
237const INSTANCE_NAME: &str = "Zed";
238
239const BUILTIN_PROFILES: &[&str] = &[
240    "CPU.Verbose.Memory",
241    "GPU.Light.Memory",
242    "DiskIO.Light.Memory",
243    "FileIO.Light.Memory",
244];
245
246fn heap_tracing_profile(heap_pid: Option<u32>) -> String {
247    let (heap_provider, heap_collector) = match heap_pid {
248        Some(pid) => (
249            format!(
250                r#"
251    <HeapEventProvider Id="ZedHeapProvider">
252      <HeapProcessIds Operation="Set">
253        <HeapProcessId Value="{pid}"/>
254      </HeapProcessIds>
255    </HeapEventProvider>"#
256            ),
257            r#"
258      <Collectors Operation="Add">
259        <HeapEventCollectorId Value="HeapCollector_WPRHeapCollector">
260          <HeapEventProviders Operation="Set">
261            <HeapEventProviderId Value="ZedHeapProvider"/>
262          </HeapEventProviders>
263        </HeapEventCollectorId>
264      </Collectors>"#
265                .to_string(),
266        ),
267        None => (String::new(), String::new()),
268    };
269
270    format!(
271        r#"<?xml version="1.0" encoding="utf-8"?>
272<WindowsPerformanceRecorder Version="1.0" Author="Zed Industries">
273  <Profiles>
274    {heap_provider}
275
276    <Profile Id="ZedHeap.Verbose.Memory" Base="Heap.Verbose.Memory" Name="ZedHeap" DetailLevel="Verbose" LoggingMode="Memory" Description="Heap tracing">
277      {heap_collector}
278    </Profile>
279  </Profiles>
280
281  <TraceMergeProperties>
282    <TraceMergeProperty Id="TraceMerge_Default" Name="TraceMerge_Default">
283      <FileCompression Value="true"/>
284    </TraceMergeProperty>
285  </TraceMergeProperties>
286</WindowsPerformanceRecorder>"#
287    )
288}
289
290fn wpr_error_context(hresult: windows_core::HRESULT, source: &windows_core::IUnknown) -> String {
291    let mut out = format!("HRESULT: {hresult}");
292
293    unsafe {
294        let mut message = BSTR::new();
295        let mut description = BSTR::new();
296        let mut detail = BSTR::new();
297        if WPRCFormatError(
298            hresult,
299            Some(source),
300            &mut message,
301            Some(&mut description),
302            Some(&mut detail),
303        )
304        .is_ok()
305        {
306            for (label, value) in [
307                ("Message", &message),
308                ("Description", &description),
309                ("Detail", &detail),
310            ] {
311                if !value.is_empty() {
312                    let _ = write!(out, "\n  {label}: {value}");
313                }
314            }
315        }
316    }
317
318    if let Ok(info) = source.cast::<IParsingErrorInfo>() {
319        unsafe {
320            if let Ok(line) = info.GetLineNumber() {
321                let _ = write!(out, "\n  Parse error at line: {line}");
322                if let Ok(col) = info.GetColumnNumber() {
323                    let _ = write!(out, ", column: {col}");
324                }
325            }
326            for (label, getter) in [
327                ("Element type", info.GetElementType()),
328                ("Element ID", info.GetElementId()),
329                ("Description", info.GetDescription()),
330            ] {
331                if let Ok(value) = getter
332                    && !value.is_empty()
333                {
334                    let _ = write!(out, "\n  {label}: {value}");
335                }
336            }
337        }
338    }
339
340    fn append_control_chain(out: &mut String, source: &windows_core::IUnknown) {
341        let Ok(info) = source.cast::<IControlErrorInfo>() else {
342            return;
343        };
344        unsafe {
345            if let Ok(object_type) = info.GetObjectType() {
346                let name = match object_type {
347                    wprcontrol::ObjectType_Profile => "Profile",
348                    wprcontrol::ObjectType_Collector => "Collector",
349                    wprcontrol::ObjectType_Provider => "Provider",
350                    _ => "Unknown",
351                };
352                let _ = write!(out, "\n  Object type: {name}");
353            }
354            if let Ok(hr) = info.GetHResult() {
355                let _ = write!(out, "\n  Inner HRESULT: {hr}");
356            }
357            if let Ok(desc) = info.GetDescription()
358                && !desc.is_empty()
359            {
360                let _ = write!(out, "\n  Description: {desc}");
361            }
362            let mut inner = None;
363            if info.GetInnerErrorInfo(&mut inner).is_ok()
364                && let Some(inner) = inner
365            {
366                let _ = write!(out, "\n  Caused by:");
367                append_control_chain(out, &inner);
368            }
369        }
370    }
371    append_control_chain(&mut out, source);
372
373    if let Ok(info) = source.cast::<windows::Win32::System::Com::IErrorInfo>() {
374        unsafe {
375            if let Ok(desc) = info.GetDescription()
376                && !desc.is_empty()
377            {
378                let _ = write!(out, "\n  IErrorInfo: {desc}");
379            }
380        }
381    }
382
383    out
384}
385
386trait WprContext<T> {
387    fn wpr_context(self, source: &impl Interface) -> Result<T>;
388}
389
390impl<T> WprContext<T> for windows_core::Result<T> {
391    fn wpr_context(self, source: &impl Interface) -> Result<T> {
392        self.map_err(|e| {
393            let unknown: windows_core::IUnknown = source.cast().expect("cast to IUnknown");
394            let context = wpr_error_context(e.code(), &unknown);
395            anyhow::anyhow!("{context}")
396        })
397    }
398}
399
400fn create_wpr<T: windows_core::Interface>(clsid: &windows_core::GUID) -> Result<T> {
401    unsafe {
402        WPRCCreateInstanceUnderInstanceName::<_, T>(
403            &BSTR::from(INSTANCE_NAME),
404            clsid,
405            None,
406            CLSCTX_INPROC_SERVER.0,
407        )
408        .context("WPRCCreateInstance failed")
409    }
410}
411
412fn build_profile_collection(heap_pid: Option<u32>) -> Result<IProfileCollection> {
413    let collection: IProfileCollection = create_wpr(&CProfileCollection)?;
414
415    for profile_name in BUILTIN_PROFILES {
416        let profile: IProfile = create_wpr(&CProfile)?;
417        unsafe {
418            profile
419                .LoadFromFile(&BSTR::from(*profile_name), &BSTR::new())
420                .wpr_context(&profile)
421                .with_context(|| format!("Load built-in profile '{profile_name}'"))?;
422            collection
423                .Add(&profile, VARIANT_FALSE)
424                .wpr_context(&collection)
425                .with_context(|| format!("Add profile '{profile_name}' to collection"))?;
426        }
427    }
428
429    let heap_xml = heap_tracing_profile(heap_pid);
430    let heap_profile: IProfile = create_wpr(&CProfile)?;
431    unsafe {
432        heap_profile
433            .LoadFromString(&BSTR::from(heap_xml))
434            .wpr_context(&heap_profile)
435            .context("Load profile from XML string")?;
436        collection
437            .Add(&heap_profile, VARIANT_BOOL(0))
438            .wpr_context(&collection)
439            .context("Add ZedHeap profile to collection")?;
440    }
441
442    Ok(collection)
443}
444
445pub fn record_etw_trace(
446    heap_pid: Option<u32>,
447    output_path: &Path,
448    socket_path: &str,
449) -> Result<()> {
450    unsafe {
451        CoInitializeEx(None, COINIT_MULTITHREADED)
452            .ok()
453            .context("COM initialization failed")?;
454    }
455
456    let socket_path = Path::new(socket_path);
457    let mut stream = net::UnixStream::connect(socket_path).context("Connect to parent socket")?;
458
459    match record_etw_trace_inner(heap_pid, output_path, &mut stream) {
460        Ok(()) => Ok(()),
461        Err(e) => {
462            send_json(
463                &mut stream,
464                &StatusMessage::Error {
465                    message: format!("{e:#}"),
466                },
467            )
468            .log_err();
469            Err(e)
470        }
471    }
472}
473
474fn record_etw_trace_inner(
475    heap_pid: Option<u32>,
476    output_path: &Path,
477    stream: &mut net::UnixStream,
478) -> Result<()> {
479    let collection = build_profile_collection(heap_pid)?;
480    let control_manager: IControlManager = create_wpr(&CControlManager)?;
481
482    // Cancel any leftover sessions with the same name that might exist
483    unsafe {
484        _ = control_manager.Cancel(None);
485    }
486
487    unsafe {
488        control_manager
489            .Start(&collection)
490            .wpr_context(&control_manager)
491            .context("Start WPR recording")?;
492    }
493
494    // We must call Save or Cancel before returning or we'll leak the kernel buffers used to record the ETW session.
495    let cancel_guard = defer({
496        let control_manager = control_manager.clone();
497        move || unsafe {
498            let _ = control_manager.Cancel(None);
499        }
500    });
501
502    send_json(stream, &StatusMessage::Started)?;
503
504    let (command, timed_out) = receive_command(stream)?;
505
506    match command {
507        Command::Cancel => {
508            unsafe {
509                control_manager
510                    .Cancel(None)
511                    .wpr_context(&control_manager)
512                    .context("Cancel WPR recording")?;
513            }
514            cancel_guard.abort();
515
516            send_json(stream, &StatusMessage::Cancelled).log_err();
517        }
518        Command::Save => {
519            unsafe {
520                control_manager
521                    .Save(
522                        &BSTR::from(output_path.to_string_lossy().as_ref()),
523                        &collection,
524                        None,
525                    )
526                    .wpr_context(&control_manager)
527                    .context("Stop WPR recording")?;
528            }
529            cancel_guard.abort();
530
531            if timed_out {
532                send_json(stream, &StatusMessage::TimedOut).log_err();
533            } else {
534                send_json(stream, &StatusMessage::Stopped).log_err();
535            }
536        }
537    }
538
539    Ok(())
540}
541
542fn receive_command(stream: &mut net::UnixStream) -> Result<(Command, bool)> {
543    use std::os::windows::io::{AsRawSocket, AsSocket};
544    use windows::Win32::Networking::WinSock::{SO_RCVTIMEO, SOL_SOCKET, setsockopt};
545
546    // Set a receive timeout so read_line returns an error after `timeout`.
547    let millis = RECORDING_TIMEOUT.as_millis() as u32;
548    let socket = stream.as_socket();
549    let ret = unsafe {
550        setsockopt(
551            windows::Win32::Networking::WinSock::SOCKET(socket.as_raw_socket() as _),
552            SOL_SOCKET,
553            SO_RCVTIMEO,
554            Some(&millis.to_ne_bytes()),
555        )
556    };
557    if ret != 0 {
558        bail!("Failed to set socket receive timeout: setsockopt returned {ret}");
559    }
560
561    let mut reader = BufReader::new(&mut *stream);
562    match recv_json::<Command>(&mut reader) {
563        Ok(command) => Ok((command, false)),
564        Err(error) => {
565            log::warn!("Failed to receive ETW command, treating as timed-out Save: {error:#}");
566            Ok((Command::Save, true))
567        }
568    }
569}
570
571pub struct EtwSession {
572    output_path: PathBuf,
573    stream: BufReader<net::UnixStream>,
574    listener: net::UnixListener,
575    socket_path: PathBuf,
576}
577
578pub fn launch_etw_recording(heap_pid: Option<u32>, output_path: &Path) -> Result<EtwSession> {
579    let sock_path = std::env::temp_dir().join(format!("zed-etw-{}.sock", std::process::id()));
580
581    _ = std::fs::remove_file(&sock_path);
582    let listener = net::UnixListener::bind(&sock_path).context("Bind Unix socket for ETW IPC")?;
583
584    let exe_path = std::env::current_exe().context("Failed to get current exe path")?;
585    let pid_arg = heap_pid.map_or(-1i64, |pid| pid as i64);
586    let args = format!(
587        "--record-etw-trace --etw-zed-pid {} --etw-output \"{}\" --etw-socket \"{}\"",
588        pid_arg,
589        output_path.display(),
590        sock_path.display(),
591    );
592
593    use windows::Win32::UI::Shell::ShellExecuteW;
594    use windows_core::PCWSTR;
595
596    let operation: Vec<u16> = "runas\0".encode_utf16().collect();
597    let file: Vec<u16> = format!("{}\0", exe_path.to_string_lossy())
598        .encode_utf16()
599        .collect();
600    let parameters: Vec<u16> = format!("{args}\0").encode_utf16().collect();
601
602    let result = unsafe {
603        ShellExecuteW(
604            None,
605            PCWSTR(operation.as_ptr()),
606            PCWSTR(file.as_ptr()),
607            PCWSTR(parameters.as_ptr()),
608            PCWSTR::null(),
609            windows::Win32::UI::WindowsAndMessaging::SW_HIDE,
610        )
611    };
612
613    let result_code = result.0 as usize;
614    if result_code <= 32 {
615        bail!("ShellExecuteW failed to launch elevated process (code: {result_code})");
616    }
617
618    let (stream, _) = listener.accept().context("Accept subprocess connection")?;
619
620    let mut session = EtwSession {
621        output_path: output_path.to_path_buf(),
622        stream: BufReader::new(stream),
623        listener,
624        socket_path: sock_path,
625    };
626
627    let status: StatusMessage =
628        recv_json(&mut session.stream).context("Wait for Started status")?;
629
630    match status {
631        StatusMessage::Started => {}
632        StatusMessage::Error { message } => {
633            bail!("Subprocess reported error during start: {message}");
634        }
635        other => {
636            bail!("Unexpected status from subprocess: {other:?}");
637        }
638    }
639
640    Ok(session)
641}
642
643#[derive(Debug, serde::Serialize, serde::Deserialize)]
644#[serde(tag = "type")]
645pub enum StatusMessage {
646    Started,
647    Stopped,
648    TimedOut,
649    Cancelled,
650    Error { message: String },
651}
652
653#[derive(Debug, serde::Serialize, serde::Deserialize)]
654#[serde(tag = "type")]
655pub enum Command {
656    Save,
657    Cancel,
658}
659
660fn send_json<T: serde::Serialize>(writer: &mut impl Write, value: &T) -> Result<()> {
661    let json = serde_json::to_string(value).context("Serialize message")?;
662    writeln!(writer, "{json}").context("Write to socket")?;
663    writer.flush().context("Flush socket")?;
664    Ok(())
665}
666
667fn recv_json<T: serde::de::DeserializeOwned>(reader: &mut impl BufRead) -> Result<T> {
668    let mut line = String::new();
669    reader.read_line(&mut line).context("Read from socket")?;
670    if line.is_empty() {
671        bail!("Socket closed before a message was received");
672    }
673    serde_json::from_str(line.trim()).context("Parse message")
674}
675