tenant.openagents/omega

1use std::{sync::Arc, thread::JoinHandle};
2
3use anyhow::Context;
4use cli::{CliRequest, CliResponse, IpcHandshake, ipc::IpcOneShotServer};
5use parking_lot::Mutex;
6use release_channel::app_identifier;
7use util::ResultExt;
8use windows::{
9    Win32::{
10        Foundation::{CloseHandle, ERROR_ALREADY_EXISTS, GENERIC_WRITE, GetLastError, HANDLE},
11        Storage::FileSystem::{
12            CreateFileW, FILE_FLAGS_AND_ATTRIBUTES, FILE_SHARE_MODE, OPEN_EXISTING,
13            PIPE_ACCESS_INBOUND, ReadFile, WriteFile,
14        },
15        System::{
16            Pipes::{
17                ConnectNamedPipe, CreateNamedPipeW, DisconnectNamedPipe, PIPE_READMODE_MESSAGE,
18                PIPE_TYPE_MESSAGE, PIPE_WAIT,
19            },
20            Threading::CreateMutexW,
21        },
22    },
23    core::HSTRING,
24};
25
26use crate::{Args, OpenListener, RawOpenRequest};
27
28#[inline]
29fn is_first_instance() -> bool {
30    unsafe {
31        CreateMutexW(
32            None,
33            false,
34            &HSTRING::from(format!("{}-Instance-Mutex", app_identifier())),
35        )
36        .expect("Unable to create instance mutex.")
37    };
38    unsafe { GetLastError() != ERROR_ALREADY_EXISTS }
39}
40
41pub fn handle_single_instance(opener: OpenListener, args: &Args) -> bool {
42    let is_first_instance = is_first_instance();
43    if is_first_instance {
44        // We are the first instance, listen for messages sent from other instances
45        std::thread::Builder::new()
46            .name("EnsureSingleton".to_owned())
47            .spawn(move || {
48                with_pipe(&|url| {
49                    opener.open(RawOpenRequest {
50                        urls: vec![url],
51                        ..Default::default()
52                    })
53                })
54            })
55            .unwrap();
56    } else if !args.foreground {
57        // We are not the first instance, send args to the first instance
58        send_args_to_instance(args).log_err();
59    }
60
61    is_first_instance
62}
63
64fn with_pipe(f: &dyn Fn(String)) {
65    let pipe = unsafe {
66        CreateNamedPipeW(
67            &HSTRING::from(format!("\\\\.\\pipe\\{}-Named-Pipe", app_identifier())),
68            PIPE_ACCESS_INBOUND,
69            PIPE_TYPE_MESSAGE | PIPE_READMODE_MESSAGE | PIPE_WAIT,
70            1,
71            128,
72            128,
73            0,
74            None,
75        )
76    };
77    if pipe.is_invalid() {
78        log::error!("Failed to create named pipe: {:?}", unsafe {
79            GetLastError()
80        });
81        return;
82    }
83
84    loop {
85        if let Some(message) = retrieve_message_from_pipe(pipe)
86            .context("Failed to read from named pipe")
87            .log_err()
88        {
89            f(message);
90        }
91    }
92}
93
94fn retrieve_message_from_pipe(pipe: HANDLE) -> anyhow::Result<String> {
95    unsafe { ConnectNamedPipe(pipe, None)? };
96    let message = retrieve_message_from_pipe_inner(pipe);
97    unsafe { DisconnectNamedPipe(pipe).log_err() };
98    message
99}
100
101fn retrieve_message_from_pipe_inner(pipe: HANDLE) -> anyhow::Result<String> {
102    let mut buffer = [0u8; 128];
103    unsafe {
104        ReadFile(pipe, Some(&mut buffer), None, None)?;
105    }
106    let message = std::ffi::CStr::from_bytes_until_nul(&buffer)?;
107    Ok(message.to_string_lossy().into_owned())
108}
109
110// This part of code is mostly from crates/cli/src/main.rs
111fn send_args_to_instance(args: &Args) -> anyhow::Result<()> {
112    if let Some(dock_menu_action_idx) = args.dock_action {
113        let url = format!("zed-dock-action://{}", dock_menu_action_idx);
114        return write_message_to_instance_pipe(url.as_bytes());
115    }
116
117    let (server, server_name) =
118        IpcOneShotServer::<IpcHandshake>::new().context("Handshake before Omega spawn")?;
119    let url = format!("zed-cli://{server_name}");
120
121    let request = {
122        let mut paths = vec![];
123        let mut urls = vec![];
124        let mut diff_paths = vec![];
125        for path in args.paths_or_urls.iter() {
126            match std::fs::canonicalize(&path) {
127                Ok(path) => paths.push(path.to_string_lossy().into_owned()),
128                Err(error) => {
129                    if path.starts_with("zed://")
130                        || path.starts_with("http://")
131                        || path.starts_with("https://")
132                        || path.starts_with("file://")
133                        || path.starts_with("ssh://")
134                    {
135                        urls.push(path.clone());
136                    } else {
137                        log::error!("error parsing path argument: {}", error);
138                    }
139                }
140            }
141        }
142
143        for path in args.diff.chunks(2) {
144            let old = std::fs::canonicalize(&path[0]).log_err();
145            let new = std::fs::canonicalize(&path[1]).log_err();
146            if let Some((old, new)) = old.zip(new) {
147                diff_paths.push([
148                    old.to_string_lossy().into_owned(),
149                    new.to_string_lossy().into_owned(),
150                ]);
151            }
152        }
153
154        CliRequest::Open {
155            paths,
156            urls,
157            diff_paths,
158            diff_all: false,
159            wait: false,
160            wsl: args.wsl.clone(),
161            open_behavior: Default::default(),
162            env: None,
163            user_data_dir: args.user_data_dir.clone(),
164            dev_container: args.dev_container,
165            cwd: std::env::current_dir().ok(),
166        }
167    };
168
169    let exit_status = Arc::new(Mutex::new(None));
170    let sender: JoinHandle<anyhow::Result<()>> = std::thread::Builder::new()
171        .name("CliReceiver".to_owned())
172        .spawn({
173            let exit_status = exit_status.clone();
174            move || {
175                let (_, handshake) = server.accept().context("Handshake after Omega spawn")?;
176                let (tx, rx) = (handshake.requests, handshake.responses);
177
178                tx.send(request)?;
179
180                while let Ok(response) = rx.recv() {
181                    match response {
182                        CliResponse::Ping => {}
183                        CliResponse::Stdout { message } => log::info!("{message}"),
184                        CliResponse::Stderr { message } => log::error!("{message}"),
185                        CliResponse::Exit { status } => {
186                            exit_status.lock().replace(status);
187                            return Ok(());
188                        }
189                        CliResponse::PromptOpenBehavior => {
190                            tx.send(CliRequest::SetOpenBehavior {
191                                behavior: cli::CliBehaviorSetting::ExistingWindow,
192                            })?;
193                        }
194                    }
195                }
196                Ok(())
197            }
198        })
199        .unwrap();
200
201    write_message_to_instance_pipe(url.as_bytes())?;
202    sender.join().unwrap()?;
203    if let Some(exit_status) = exit_status.lock().take() {
204        std::process::exit(exit_status);
205    }
206    Ok(())
207}
208
209fn write_message_to_instance_pipe(message: &[u8]) -> anyhow::Result<()> {
210    unsafe {
211        let pipe = CreateFileW(
212            &HSTRING::from(format!("\\\\.\\pipe\\{}-Named-Pipe", app_identifier())),
213            GENERIC_WRITE.0,
214            FILE_SHARE_MODE::default(),
215            None,
216            OPEN_EXISTING,
217            FILE_FLAGS_AND_ATTRIBUTES::default(),
218            None,
219        )?;
220        WriteFile(pipe, Some(message), None, None)?;
221        CloseHandle(pipe)?;
222    }
223    Ok(())
224}
225