1use std::path::PathBuf;
  2
  3use anyhow::Context;
  4use collections::HashMap;
  5use schemars::{gen::SchemaSettings, JsonSchema};
  6use serde::{Deserialize, Serialize};
  7use sha2::{Digest, Sha256};
  8use util::{truncate_and_remove_front, ResultExt};
  9
 10use crate::{ResolvedTask, SpawnInTerminal, TaskContext, TaskId, ZED_VARIABLE_NAME_PREFIX};
 11
 12/// A template definition of a Zed task to run.
 13/// May use the [`VariableName`] to get the corresponding substitutions into its fields.
 14///
 15/// Template itself is not ready to spawn a task, it needs to be resolved with a [`TaskContext`] first, that
 16/// contains all relevant Zed state in task variables.
 17/// A single template may produce different tasks (or none) for different contexts.
 18#[derive(Clone, Default, Debug, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
 19#[serde(rename_all = "snake_case")]
 20pub struct TaskTemplate {
 21    /// Human readable name of the task to display in the UI.
 22    pub label: String,
 23    /// Executable command to spawn.
 24    pub command: String,
 25    /// Arguments to the command.
 26    #[serde(default)]
 27    pub args: Vec<String>,
 28    /// Env overrides for the command, will be appended to the terminal's environment from the settings.
 29    #[serde(default)]
 30    pub env: HashMap<String, String>,
 31    /// Current working directory to spawn the command into, defaults to current project root.
 32    #[serde(default)]
 33    pub cwd: Option<String>,
 34    /// Whether to use a new terminal tab or reuse the existing one to spawn the process.
 35    #[serde(default)]
 36    pub use_new_terminal: bool,
 37    /// Whether to allow multiple instances of the same task to be run, or rather wait for the existing ones to finish.
 38    #[serde(default)]
 39    pub allow_concurrent_runs: bool,
 40    /// What to do with the terminal pane and tab, after the command was started:
 41    /// * `always` — always show the terminal pane, add and focus the corresponding task's tab in it (default)
 42    /// * `never` — avoid changing current terminal pane focus, but still add/reuse the task's tab there
 43    #[serde(default)]
 44    pub reveal: RevealStrategy,
 45}
 46
 47/// What to do with the terminal pane and tab, after the command was started.
 48#[derive(Default, Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
 49#[serde(rename_all = "snake_case")]
 50pub enum RevealStrategy {
 51    /// Always show the terminal pane, add and focus the corresponding task's tab in it.
 52    #[default]
 53    Always,
 54    /// Do not change terminal pane focus, but still add/reuse the task's tab there.
 55    Never,
 56}
 57
 58/// A group of Tasks defined in a JSON file.
 59#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
 60pub struct TaskTemplates(pub Vec<TaskTemplate>);
 61
 62impl TaskTemplates {
 63    /// Generates JSON schema of Tasks JSON template format.
 64    pub fn generate_json_schema() -> serde_json_lenient::Value {
 65        let schema = SchemaSettings::draft07()
 66            .with(|settings| settings.option_add_null_type = false)
 67            .into_generator()
 68            .into_root_schema_for::<Self>();
 69
 70        serde_json_lenient::to_value(schema).unwrap()
 71    }
 72}
 73
 74impl TaskTemplate {
 75    /// Replaces all `VariableName` task variables in the task template string fields.
 76    /// If any replacement fails or the new string substitutions still have [`ZED_VARIABLE_NAME_PREFIX`],
 77    /// `None` is returned.
 78    ///
 79    /// Every [`ResolvedTask`] gets a [`TaskId`], based on the `id_base` (to avoid collision with various task sources),
 80    /// and hashes of its template and [`TaskContext`], see [`ResolvedTask`] fields' documentation for more details.
 81    pub fn resolve_task(&self, id_base: &str, cx: TaskContext) -> Option<ResolvedTask> {
 82        if self.label.trim().is_empty() || self.command.trim().is_empty() {
 83            return None;
 84        }
 85        let TaskContext {
 86            cwd,
 87            task_variables,
 88        } = cx;
 89        let task_variables = task_variables.into_env_variables();
 90        let truncated_variables = truncate_variables(&task_variables);
 91        let cwd = match self.cwd.as_deref() {
 92            Some(cwd) => Some(substitute_all_template_variables_in_str(
 93                cwd,
 94                &task_variables,
 95            )?),
 96            None => None,
 97        }
 98        .map(PathBuf::from)
 99        .or(cwd);
100        let shortened_label =
101            substitute_all_template_variables_in_str(&self.label, &truncated_variables)?;
102        let full_label = substitute_all_template_variables_in_str(&self.label, &task_variables)?;
103        let command = substitute_all_template_variables_in_str(&self.command, &task_variables)?;
104        let args = substitute_all_template_variables_in_vec(self.args.clone(), &task_variables)?;
105        let task_hash = to_hex_hash(self)
106            .context("hashing task template")
107            .log_err()?;
108        let variables_hash = to_hex_hash(&task_variables)
109            .context("hashing task variables")
110            .log_err()?;
111        let id = TaskId(format!("{id_base}_{task_hash}_{variables_hash}"));
112        let mut env = substitute_all_template_variables_in_map(self.env.clone(), &task_variables)?;
113        env.extend(task_variables);
114        Some(ResolvedTask {
115            id: id.clone(),
116            original_task: self.clone(),
117            resolved_label: full_label,
118            resolved: Some(SpawnInTerminal {
119                id,
120                cwd,
121                label: shortened_label,
122                command,
123                args,
124                env,
125                use_new_terminal: self.use_new_terminal,
126                allow_concurrent_runs: self.allow_concurrent_runs,
127                reveal: self.reveal,
128            }),
129        })
130    }
131}
132
133const MAX_DISPLAY_VARIABLE_LENGTH: usize = 15;
134
135fn truncate_variables(task_variables: &HashMap<String, String>) -> HashMap<String, String> {
136    task_variables
137        .iter()
138        .map(|(key, value)| {
139            (
140                key.clone(),
141                truncate_and_remove_front(value, MAX_DISPLAY_VARIABLE_LENGTH),
142            )
143        })
144        .collect()
145}
146
147fn to_hex_hash(object: impl Serialize) -> anyhow::Result<String> {
148    let json = serde_json_lenient::to_string(&object).context("serializing the object")?;
149    let mut hasher = Sha256::new();
150    hasher.update(json.as_bytes());
151    Ok(hex::encode(hasher.finalize()))
152}
153
154fn substitute_all_template_variables_in_str(
155    template_str: &str,
156    task_variables: &HashMap<String, String>,
157) -> Option<String> {
158    let substituted_string = subst::substitute(&template_str, task_variables).ok()?;
159    if substituted_string.contains(ZED_VARIABLE_NAME_PREFIX) {
160        return None;
161    }
162    Some(substituted_string)
163}
164
165fn substitute_all_template_variables_in_vec(
166    mut template_strs: Vec<String>,
167    task_variables: &HashMap<String, String>,
168) -> Option<Vec<String>> {
169    for template_str in &mut template_strs {
170        let substituted_string = subst::substitute(&template_str, task_variables).ok()?;
171        if substituted_string.contains(ZED_VARIABLE_NAME_PREFIX) {
172            return None;
173        }
174        *template_str = substituted_string
175    }
176    Some(template_strs)
177}
178
179fn substitute_all_template_variables_in_map(
180    keys_and_values: HashMap<String, String>,
181    task_variables: &HashMap<String, String>,
182) -> Option<HashMap<String, String>> {
183    keys_and_values
184        .into_iter()
185        .try_fold(HashMap::default(), |mut expanded_keys, (mut key, value)| {
186            match task_variables.get(&key) {
187                Some(variable_expansion) => key = variable_expansion.clone(),
188                None => {
189                    if key.starts_with(ZED_VARIABLE_NAME_PREFIX) {
190                        return Err(());
191                    }
192                }
193            }
194            expanded_keys.insert(
195                key,
196                subst::substitute(&value, task_variables)
197                    .map_err(|_| ())?
198                    .to_string(),
199            );
200            Ok(expanded_keys)
201        })
202        .ok()
203}
204
205#[cfg(test)]
206mod tests {
207    use std::{borrow::Cow, path::Path};
208
209    use crate::{TaskVariables, VariableName};
210
211    use super::*;
212
213    const TEST_ID_BASE: &str = "test_base";
214
215    #[test]
216    fn test_resolving_templates_with_blank_command_and_label() {
217        let task_with_all_properties = TaskTemplate {
218            label: "test_label".to_string(),
219            command: "test_command".to_string(),
220            args: vec!["test_arg".to_string()],
221            env: HashMap::from_iter([("test_env_key".to_string(), "test_env_var".to_string())]),
222            ..TaskTemplate::default()
223        };
224
225        for task_with_blank_property in &[
226            TaskTemplate {
227                label: "".to_string(),
228                ..task_with_all_properties.clone()
229            },
230            TaskTemplate {
231                command: "".to_string(),
232                ..task_with_all_properties.clone()
233            },
234            TaskTemplate {
235                label: "".to_string(),
236                command: "".to_string(),
237                ..task_with_all_properties.clone()
238            },
239        ] {
240            assert_eq!(
241                task_with_blank_property.resolve_task(TEST_ID_BASE, TaskContext::default()),
242                None,
243                "should not resolve task with blank label and/or command: {task_with_blank_property:?}"
244            );
245        }
246    }
247
248    #[test]
249    fn test_template_cwd_resolution() {
250        let task_without_cwd = TaskTemplate {
251            cwd: None,
252            label: "test task".to_string(),
253            command: "echo 4".to_string(),
254            ..TaskTemplate::default()
255        };
256
257        let resolved_task = |task_template: &TaskTemplate, task_cx| {
258            let resolved_task = task_template
259                .resolve_task(TEST_ID_BASE, task_cx)
260                .unwrap_or_else(|| panic!("failed to resolve task {task_without_cwd:?}"));
261            resolved_task
262                .resolved
263                .clone()
264                .unwrap_or_else(|| {
265                    panic!("failed to get resolve data for resolved task. Template: {task_without_cwd:?} Resolved: {resolved_task:?}")
266                })
267        };
268
269        assert_eq!(
270            resolved_task(
271                &task_without_cwd,
272                TaskContext {
273                    cwd: None,
274                    task_variables: TaskVariables::default(),
275                }
276            )
277            .cwd,
278            None,
279            "When neither task nor task context have cwd, it should be None"
280        );
281
282        let context_cwd = Path::new("a").join("b").join("c");
283        assert_eq!(
284            resolved_task(
285                &task_without_cwd,
286                TaskContext {
287                    cwd: Some(context_cwd.clone()),
288                    task_variables: TaskVariables::default(),
289                }
290            )
291            .cwd
292            .as_deref(),
293            Some(context_cwd.as_path()),
294            "TaskContext's cwd should be taken on resolve if task's cwd is None"
295        );
296
297        let task_cwd = Path::new("d").join("e").join("f");
298        let mut task_with_cwd = task_without_cwd.clone();
299        task_with_cwd.cwd = Some(task_cwd.display().to_string());
300        let task_with_cwd = task_with_cwd;
301
302        assert_eq!(
303            resolved_task(
304                &task_with_cwd,
305                TaskContext {
306                    cwd: None,
307                    task_variables: TaskVariables::default(),
308                }
309            )
310            .cwd
311            .as_deref(),
312            Some(task_cwd.as_path()),
313            "TaskTemplate's cwd should be taken on resolve if TaskContext's cwd is None"
314        );
315
316        assert_eq!(
317            resolved_task(
318                &task_with_cwd,
319                TaskContext {
320                    cwd: Some(context_cwd.clone()),
321                    task_variables: TaskVariables::default(),
322                }
323            )
324            .cwd
325            .as_deref(),
326            Some(task_cwd.as_path()),
327            "TaskTemplate's cwd should be taken on resolve if TaskContext's cwd is not None"
328        );
329    }
330
331    #[test]
332    fn test_template_variables_resolution() {
333        let custom_variable_1 = VariableName::Custom(Cow::Borrowed("custom_variable_1"));
334        let custom_variable_2 = VariableName::Custom(Cow::Borrowed("custom_variable_2"));
335        let long_value = "01".repeat(MAX_DISPLAY_VARIABLE_LENGTH * 2);
336        let all_variables = [
337            (VariableName::Row, "1234".to_string()),
338            (VariableName::Column, "5678".to_string()),
339            (VariableName::File, "test_file".to_string()),
340            (VariableName::SelectedText, "test_selected_text".to_string()),
341            (VariableName::Symbol, long_value.clone()),
342            (VariableName::WorktreeRoot, "/test_root/".to_string()),
343            (
344                custom_variable_1.clone(),
345                "test_custom_variable_1".to_string(),
346            ),
347            (
348                custom_variable_2.clone(),
349                "test_custom_variable_2".to_string(),
350            ),
351        ];
352
353        let task_with_all_variables = TaskTemplate {
354            label: format!(
355                "test label for {} and {}",
356                VariableName::Row.template_value(),
357                VariableName::Symbol.template_value(),
358            ),
359            command: format!(
360                "echo {} {}",
361                VariableName::File.template_value(),
362                VariableName::Symbol.template_value(),
363            ),
364            args: vec![
365                format!("arg1 {}", VariableName::SelectedText.template_value()),
366                format!("arg2 {}", VariableName::Column.template_value()),
367                format!("arg3 {}", VariableName::Symbol.template_value()),
368            ],
369            env: HashMap::from_iter([
370                ("test_env_key".to_string(), "test_env_var".to_string()),
371                (
372                    "env_key_1".to_string(),
373                    VariableName::WorktreeRoot.template_value(),
374                ),
375                (
376                    "env_key_2".to_string(),
377                    format!(
378                        "env_var_2_{}_{}",
379                        custom_variable_1.template_value(),
380                        custom_variable_2.template_value()
381                    ),
382                ),
383                (
384                    "env_key_3".to_string(),
385                    format!("env_var_3_{}", VariableName::Symbol.template_value()),
386                ),
387            ]),
388            ..TaskTemplate::default()
389        };
390
391        let mut first_resolved_id = None;
392        for i in 0..15 {
393            let resolved_task = task_with_all_variables.resolve_task(
394                TEST_ID_BASE,
395                TaskContext {
396                    cwd: None,
397                    task_variables: TaskVariables::from_iter(all_variables.clone()),
398                },
399            ).unwrap_or_else(|| panic!("Should successfully resolve task {task_with_all_variables:?} with variables {all_variables:?}"));
400
401            match &first_resolved_id {
402                None => first_resolved_id = Some(resolved_task.id),
403                Some(first_id) => assert_eq!(
404                    &resolved_task.id, first_id,
405                    "Step {i}, for the same task template and context, there should be the same resolved task id"
406                ),
407            }
408
409            assert_eq!(
410                resolved_task.original_task, task_with_all_variables,
411                "Resolved task should store its template without changes"
412            );
413            assert_eq!(
414                resolved_task.resolved_label,
415                format!("test label for 1234 and {long_value}"),
416                "Resolved task label should be substituted with variables and those should not be shortened"
417            );
418
419            let spawn_in_terminal = resolved_task
420                .resolved
421                .as_ref()
422                .expect("should have resolved a spawn in terminal task");
423            assert_eq!(
424                spawn_in_terminal.label,
425                format!(
426                    "test label for 1234 and …{}",
427                    &long_value[..=MAX_DISPLAY_VARIABLE_LENGTH]
428                ),
429                "Human-readable label should have long substitutions trimmed"
430            );
431            assert_eq!(
432                spawn_in_terminal.command,
433                format!("echo test_file {long_value}"),
434                "Command should be substituted with variables and those should not be shortened"
435            );
436            assert_eq!(
437                spawn_in_terminal.args,
438                &[
439                    "arg1 test_selected_text",
440                    "arg2 5678",
441                    &format!("arg3 {long_value}")
442                ],
443                "Args should be substituted with variables and those should not be shortened"
444            );
445
446            assert_eq!(
447                spawn_in_terminal
448                    .env
449                    .get("test_env_key")
450                    .map(|s| s.as_str()),
451                Some("test_env_var")
452            );
453            assert_eq!(
454                spawn_in_terminal.env.get("env_key_1").map(|s| s.as_str()),
455                Some("/test_root/")
456            );
457            assert_eq!(
458                spawn_in_terminal.env.get("env_key_2").map(|s| s.as_str()),
459                Some("env_var_2_test_custom_variable_1_test_custom_variable_2")
460            );
461            assert_eq!(
462                spawn_in_terminal.env.get("env_key_3"),
463                Some(&format!("env_var_3_{long_value}")),
464                "Env vars should be substituted with variables and those should not be shortened"
465            );
466        }
467
468        for i in 0..all_variables.len() {
469            let mut not_all_variables = all_variables.to_vec();
470            let removed_variable = not_all_variables.remove(i);
471            let resolved_task_attempt = task_with_all_variables.resolve_task(
472                TEST_ID_BASE,
473                TaskContext {
474                    cwd: None,
475                    task_variables: TaskVariables::from_iter(not_all_variables),
476                },
477            );
478            assert_eq!(resolved_task_attempt, None, "If any of the Zed task variables is not substituted, the task should not be resolved, but got some resolution without the variable {removed_variable:?} (index {i})");
479        }
480    }
481}