tenant.openagents/omega

1use anyhow::{Context as _, Result};
2use async_trait::async_trait;
3use dap::{DapLocator, DebugRequest, adapters::DebugAdapterName};
4use gpui::{BackgroundExecutor, SharedString};
5use serde_json::{Value, json};
6use smol::{io::AsyncReadExt, process::Stdio as SmolStdio};
7use std::time::Duration;
8use task::{BuildTaskDefinition, DebugScenario, ShellBuilder, SpawnInTerminal, TaskTemplate};
9use util::command::{Stdio, new_command};
10
11pub(crate) struct CargoLocator;
12
13async fn find_best_executable(
14    executables: &[String],
15    test_name: &str,
16    executor: BackgroundExecutor,
17) -> Option<String> {
18    if executables.len() == 1 {
19        return executables.first().cloned();
20    }
21    for executable in executables {
22        let Some(mut child) = new_command(&executable)
23            .arg("--list")
24            .stdout(Stdio::piped())
25            .spawn()
26            .ok()
27        else {
28            continue;
29        };
30        let mut test_lines = String::default();
31        let exec_result = smol::future::race(
32            async {
33                if let Some(mut stdout) = child.stdout.take() {
34                    stdout.read_to_string(&mut test_lines).await?;
35                }
36                Ok(())
37            },
38            async {
39                executor.timer(Duration::from_secs(3)).await;
40                anyhow::bail!("Timed out waiting for executable stdout")
41            },
42        );
43
44        if let Err(err) = exec_result.await {
45            log::warn!("Failed to list tests for {executable}: {err}");
46        } else {
47            for line in test_lines.lines() {
48                if line.contains(&test_name) {
49                    return Some(executable.clone());
50                }
51            }
52        }
53        let _ = child.kill();
54    }
55    None
56}
57#[async_trait]
58impl DapLocator for CargoLocator {
59    fn name(&self) -> SharedString {
60        SharedString::new_static("rust-cargo-locator")
61    }
62    async fn create_scenario(
63        &self,
64        build_config: &TaskTemplate,
65        resolved_label: &str,
66        adapter: &DebugAdapterName,
67    ) -> Option<DebugScenario> {
68        if build_config.command != "cargo" {
69            return None;
70        }
71        let mut task_template = build_config.clone();
72        let cargo_action = task_template.args.first_mut()?;
73        if cargo_action == "check" || cargo_action == "clean" {
74            return None;
75        }
76
77        match cargo_action.as_ref() {
78            "run" | "r" => {
79                *cargo_action = "build".to_owned();
80            }
81            "test" | "t" | "bench" => {
82                let delimiter = task_template
83                    .args
84                    .iter()
85                    .position(|arg| arg == "--")
86                    .unwrap_or(task_template.args.len());
87                if !task_template.args[..delimiter]
88                    .iter()
89                    .any(|arg| arg == "--no-run")
90                {
91                    task_template.args.insert(delimiter, "--no-run".to_owned());
92                }
93            }
94            _ => {}
95        }
96
97        let config = if adapter.as_ref() == "CodeLLDB" {
98            json!({
99                "sourceLanguages": ["rust"]
100            })
101        } else {
102            Value::Null
103        };
104        Some(DebugScenario {
105            adapter: adapter.0.clone(),
106            label: resolved_label.to_string().into(),
107            build: Some(BuildTaskDefinition::Template {
108                task_template,
109                locator_name: Some(self.name()),
110            }),
111            config,
112            tcp_connection: None,
113        })
114    }
115
116    async fn run(
117        &self,
118        build_config: SpawnInTerminal,
119        executor: BackgroundExecutor,
120    ) -> Result<DebugRequest> {
121        let cwd = build_config
122            .cwd
123            .clone()
124            .context("Couldn't get cwd from debug config which is needed for locators")?;
125        let builder = ShellBuilder::new(&build_config.shell, cfg!(windows)).non_interactive();
126        let mut child = builder
127            .build_smol_command(
128                Some("cargo".into()),
129                &build_config
130                    .args
131                    .iter()
132                    .cloned()
133                    .take_while(|arg| arg != "--")
134                    .chain(Some("--message-format=json".to_owned()))
135                    .collect::<Vec<_>>(),
136            )
137            .envs(build_config.env.iter().map(|(k, v)| (k.clone(), v.clone())))
138            .current_dir(cwd)
139            .stdout(SmolStdio::piped())
140            .spawn()?;
141
142        let mut output = String::new();
143        if let Some(mut stdout) = child.stdout.take() {
144            stdout.read_to_string(&mut output).await?;
145        }
146
147        let status = child.status().await?;
148        anyhow::ensure!(status.success(), "Cargo command failed");
149
150        let is_test = build_config
151            .args
152            .first()
153            .is_some_and(|arg| arg == "test" || arg == "t");
154
155        let is_ignored = build_config.args.contains(&"--include-ignored".to_owned());
156
157        let executables = output
158            .lines()
159            .filter(|line| !line.trim().is_empty())
160            .filter_map(|line| serde_json::from_str(line).ok())
161            .filter(|json: &Value| {
162                let is_test_binary = json
163                    .get("profile")
164                    .and_then(|profile| profile.get("test"))
165                    .and_then(Value::as_bool)
166                    .unwrap_or(false);
167
168                if is_test {
169                    is_test_binary
170                } else {
171                    !is_test_binary
172                }
173            })
174            .filter_map(|json: Value| {
175                json.get("executable")
176                    .and_then(Value::as_str)
177                    .map(String::from)
178            })
179            .collect::<Vec<_>>();
180        anyhow::ensure!(
181            !executables.is_empty(),
182            "Couldn't get executable in cargo locator"
183        );
184
185        let mut test_name = None;
186        if is_test {
187            test_name = build_config
188                .args
189                .iter()
190                .rev()
191                .take_while(|name| "--" != name.as_str())
192                .find(|name| !name.starts_with("-"))
193                .cloned();
194        }
195        let executable = {
196            if let Some(name) = test_name.as_ref().and_then(|name| {
197                name.strip_prefix('$')
198                    .map(|name| build_config.env.get(name))
199                    .unwrap_or(Some(name))
200            }) {
201                find_best_executable(&executables, name, executor).await
202            } else {
203                None
204            }
205        };
206
207        let Some(executable) = executable.or_else(|| executables.first().cloned()) else {
208            anyhow::bail!("Couldn't get executable in cargo locator");
209        };
210
211        let args = build_test_binary_args(test_name.as_deref(), is_test, is_ignored);
212
213        Ok(DebugRequest::Launch(task::LaunchRequest {
214            program: executable,
215            cwd: build_config.cwd,
216            args,
217            env: build_config.env.into_iter().collect(),
218        }))
219    }
220}
221
222fn build_test_binary_args(test_name: Option<&str>, is_test: bool, is_ignored: bool) -> Vec<String> {
223    let mut args: Vec<String> = test_name.map(str::to_owned).into_iter().collect();
224    if is_test {
225        args.push("--nocapture".to_owned());
226        if is_ignored {
227            args.push("--include-ignored".to_owned());
228            // Append `--exact` only if we can be sure the name is fully
229            // qualified. Runnables produced from tree-sitter are not qualified
230            // (see #51810).
231            if test_name.is_some_and(|name| name.contains("::")) {
232                args.push("--exact".to_owned());
233            }
234        }
235    }
236    args
237}
238
239#[cfg(test)]
240mod tests {
241    use super::build_test_binary_args;
242
243    #[test]
244    fn non_test_invocation_has_no_test_args() {
245        assert_eq!(
246            build_test_binary_args(None, false, false),
247            Vec::<String>::new()
248        );
249    }
250
251    #[test]
252    fn bare_test_name_does_not_get_exact() {
253        // Zed's tree-sitter runnable template for Rust tests captures only the function
254        // identifier and always passes `--include-ignored`, so `is_ignored` is true here
255        // for a regular (non-ignored) test.
256        assert_eq!(
257            build_test_binary_args(Some("get_complex_variant"), true, true),
258            vec![
259                "get_complex_variant".to_owned(),
260                "--nocapture".to_owned(),
261                "--include-ignored".to_owned(),
262            ],
263        );
264    }
265
266    #[test]
267    fn qualified_test_name_gets_exact() {
268        assert_eq!(
269            build_test_binary_args(Some("variant_get::test::get_complex_variant"), true, true,),
270            vec![
271                "variant_get::test::get_complex_variant".to_owned(),
272                "--nocapture".to_owned(),
273                "--include-ignored".to_owned(),
274                "--exact".to_owned(),
275            ],
276        );
277    }
278
279    #[test]
280    fn test_without_include_ignored_never_gets_exact() {
281        assert_eq!(
282            build_test_binary_args(Some("variant_get::test::get_complex_variant"), true, false,),
283            vec![
284                "variant_get::test::get_complex_variant".to_owned(),
285                "--nocapture".to_owned(),
286            ],
287        );
288    }
289}
290