tenant.openagents/omega

1use std::{sync::Arc, time::Duration};
2
3use anyhow::anyhow;
4use askpass::EncryptedPassword;
5use clap::Parser;
6use client::{Client, UserStore};
7use futures::channel::oneshot;
8use gpui::AppContext as _;
9use gpui::TaskExt;
10use http_client::FakeHttpClient;
11use language::LanguageRegistry;
12use node_runtime::NodeRuntime;
13use project::{
14    Project, RealFs,
15    search::{SearchQuery, SearchResult},
16};
17use release_channel::ReleaseChannel;
18use remote::{ConnectionIdentifier, RemoteClientDelegate, SshConnectionOptions};
19use semver::Version;
20
21#[derive(Parser)]
22struct Args {
23    /// List of worktrees to run the search against.
24    worktrees: Vec<String>,
25    #[clap(short)]
26    query: Option<String>,
27    /// Askpass socket for SSH authentication
28    #[clap(long)]
29    askpass: Option<String>,
30    /// Treat query as a regex.
31    #[clap(short, long)]
32    regex: bool,
33    /// Matches have to be standalone words.
34    #[clap(long)]
35    whole_word: bool,
36    /// Make matching case-sensitive.
37    #[clap(long, default_value_t = false)]
38    case_sensitive: bool,
39    /// Include gitignored files in the search.
40    #[clap(long)]
41    include_ignored: bool,
42    #[clap(long)]
43    ssh: Option<String>,
44}
45
46struct BenchmarkRemoteClient;
47impl RemoteClientDelegate for BenchmarkRemoteClient {
48    fn ask_password(
49        &self,
50        prompt: String,
51        tx: oneshot::Sender<EncryptedPassword>,
52        _cx: &mut gpui::AsyncApp,
53    ) {
54        eprintln!("SSH asking for password: {}", prompt);
55        match rpassword::prompt_password(&prompt) {
56            Ok(password) => match EncryptedPassword::try_from(password.as_ref()) {
57                Ok(encrypted) => {
58                    if tx.send(encrypted).is_err() {
59                        eprintln!("Failed to send password");
60                    }
61                }
62                Err(e) => eprintln!("Failed to encrypt password: {e}"),
63            },
64            Err(e) => eprintln!("Failed to read password: {e}"),
65        }
66    }
67
68    fn get_download_url(
69        &self,
70        _platform: remote::RemotePlatform,
71        _release_channel: ReleaseChannel,
72        _version: Option<Version>,
73        _cx: &mut gpui::AsyncApp,
74    ) -> gpui::Task<gpui::Result<Option<String>>> {
75        unimplemented!()
76    }
77
78    fn download_server_binary_locally(
79        &self,
80        _platform: remote::RemotePlatform,
81        _release_channel: ReleaseChannel,
82        _version: Option<Version>,
83        _cx: &mut gpui::AsyncApp,
84    ) -> gpui::Task<gpui::Result<std::path::PathBuf>> {
85        unimplemented!()
86    }
87
88    fn set_status(&self, status: Option<&str>, _: &mut gpui::AsyncApp) {
89        if let Some(status) = status {
90            println!("SSH status: {status}");
91        }
92    }
93}
94fn main() -> Result<(), anyhow::Error> {
95    let args = Args::parse();
96
97    if let Some(socket) = &args.askpass {
98        askpass::main(socket);
99        return Ok(());
100    }
101
102    let query_str = args
103        .query
104        .ok_or_else(|| anyhow!("-q/--query is required"))?;
105    let query = if args.regex {
106        SearchQuery::regex(
107            query_str,
108            args.whole_word,
109            args.case_sensitive,
110            args.include_ignored,
111            false,
112            Default::default(),
113            Default::default(),
114            false,
115            None,
116        )
117    } else {
118        SearchQuery::text(
119            query_str,
120            args.whole_word,
121            args.case_sensitive,
122            args.include_ignored,
123            Default::default(),
124            Default::default(),
125            false,
126            None,
127        )
128    }?;
129    gpui_platform::headless().run(|cx| {
130        release_channel::init_test(semver::Version::new(0, 0, 0), ReleaseChannel::Dev, cx);
131        settings::init(cx);
132        let client = Client::production(cx);
133        let http_client = FakeHttpClient::with_200_response();
134        let (_, rx) = watch::channel(None);
135        let node = NodeRuntime::new(http_client, None, rx);
136        let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
137        let registry = Arc::new(LanguageRegistry::new(cx.background_executor().clone()));
138        let fs = Arc::new(RealFs::new(None, cx.background_executor().clone()));
139
140
141
142            cx.spawn(async move |cx| {
143                let project = if let Some(ssh_target) = args.ssh {
144                    println!("Setting up SSH connection for {}", &ssh_target);
145                    let ssh_connection_options = SshConnectionOptions::parse_command_line(&ssh_target)?;
146
147                    let connection_options = remote::RemoteConnectionOptions::from(ssh_connection_options);
148                    let delegate = Arc::new(BenchmarkRemoteClient);
149                    let remote_connection = remote::connect(connection_options.clone(), delegate.clone(), cx).await.unwrap();
150
151                    let (_tx, rx) = oneshot::channel();
152                    let remote_client =  cx.update(|cx| remote::RemoteClient::new(ConnectionIdentifier::setup(), remote_connection, rx, delegate.clone(), cx )).await?.ok_or_else(|| anyhow!("ssh initialization returned None"))?;
153
154                    cx.update(|cx| Project::remote(remote_client,  client, node, user_store, registry, fs, false, cx))
155                } else {
156                    println!("Setting up local project");
157                    cx.update(|cx| Project::local(
158                    client,
159                    node,
160                    user_store,
161                    registry,
162                    fs,
163                    Some(Default::default()),
164                    project::LocalProjectFlags {
165                        init_worktree_trust: false,
166                        ..Default::default()
167                    },
168                    cx,
169                ))
170                };
171                println!("Loading worktrees");
172                let worktrees = project.update(cx, |this, cx| {
173                    args.worktrees
174                        .into_iter()
175                        .map(|worktree| this.find_or_create_worktree(worktree, true, cx))
176                        .collect::<Vec<_>>()
177                });
178
179                let worktrees = futures::future::join_all(worktrees)
180                    .await
181                    .into_iter()
182                    .collect::<Result<Vec<_>, anyhow::Error>>()?;
183
184                for (worktree, _) in &worktrees {
185                    let scan_complete = worktree
186                        .update(cx, |this, _| {
187                            if let Some(local) = this.as_local() {
188                                Some(local.scan_complete())
189                            } else {
190                                None
191                            }
192                        });
193                    if let Some(scan_complete) = scan_complete {
194                        scan_complete.await;
195                    } else {
196                        cx.background_executor().timer(Duration::from_secs(10)).await;
197                    }
198
199                }
200                println!("Worktrees loaded");
201
202                println!("Starting a project search");
203                let timer = std::time::Instant::now();
204                let mut first_match = None;
205                let matches = project.update(cx, |this, cx| this.search(query, cx));
206                let mut matched_files = 0;
207                let mut matched_chunks = 0;
208                while let Ok(match_result) = matches.rx.recv().await {
209                    if first_match.is_none() {
210                        let time = timer.elapsed();
211                        first_match = Some(time);
212                        println!("First match found after {time:?}");
213                    }
214                    match match_result {
215                        SearchResult::Buffer { ranges, .. } => {
216                            matched_files += 1;
217                            matched_chunks += ranges.len();
218                        }
219                        SearchResult::LimitReached => break,
220                        SearchResult::WaitingForScan | SearchResult::Searching => continue,
221                    }
222                }
223                let elapsed = timer.elapsed();
224                println!(
225                    "Finished project search after {elapsed:?}. Matched {matched_files} files and {matched_chunks} excerpts"
226                );
227                drop(project);
228                cx.update(|cx| cx.quit());
229
230                anyhow::Ok(())
231            })
232            .detach_and_log_err(cx);
233
234    });
235    Ok(())
236}
237