task_template.rs

  1use std::path::PathBuf;
  2
  3use anyhow::{bail, 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
106        let task_hash = to_hex_hash(&self)
107            .context("hashing task template")
108            .log_err()?;
109        let variables_hash = to_hex_hash(&task_variables)
110            .context("hashing task variables")
111            .log_err()?;
112        let id = TaskId(format!("{id_base}_{task_hash}_{variables_hash}"));
113        let mut env = substitute_all_template_variables_in_map(self.env.clone(), &task_variables)?;
114        env.extend(task_variables);
115        Some(ResolvedTask {
116            id: id.clone(),
117            original_task: self.clone(),
118            resolved_label: full_label,
119            resolved: Some(SpawnInTerminal {
120                id,
121                cwd,
122                label: shortened_label,
123                command,
124                args,
125                env,
126                use_new_terminal: self.use_new_terminal,
127                allow_concurrent_runs: self.allow_concurrent_runs,
128                reveal: self.reveal,
129            }),
130        })
131    }
132}
133
134const MAX_DISPLAY_VARIABLE_LENGTH: usize = 15;
135
136fn truncate_variables(task_variables: &HashMap<String, String>) -> HashMap<String, String> {
137    task_variables
138        .iter()
139        .map(|(key, value)| {
140            (
141                key.clone(),
142                truncate_and_remove_front(value, MAX_DISPLAY_VARIABLE_LENGTH),
143            )
144        })
145        .collect()
146}
147
148fn to_hex_hash(object: impl Serialize) -> anyhow::Result<String> {
149    let json = serde_json_lenient::to_string(&object).context("serializing the object")?;
150    let mut hasher = Sha256::new();
151    hasher.update(json.as_bytes());
152    Ok(hex::encode(hasher.finalize()))
153}
154
155fn substitute_all_template_variables_in_str(
156    template_str: &str,
157    task_variables: &HashMap<String, String>,
158) -> Option<String> {
159    let substituted_string = shellexpand::env_with_context(&template_str, |var| {
160        // Colons denote a default value in case the variable is not set. We want to preserve that default, as otherwise shellexpand will substitute it for us.
161        let colon_position = var.find(':').unwrap_or(var.len());
162        let (variable_name, default) = var.split_at(colon_position);
163        let append_previous_default = |ret: &mut String| {
164            if !default.is_empty() {
165                ret.push_str(default);
166            }
167        };
168        if let Some(mut name) = task_variables.get(variable_name).cloned() {
169            // Got a task variable hit
170            append_previous_default(&mut name);
171            return Ok(Some(name));
172        } else if variable_name.starts_with(ZED_VARIABLE_NAME_PREFIX) {
173            bail!("Unknown variable name: {}", variable_name);
174        }
175        // This is an unknown variable.
176        // We should not error out, as they may come from user environment (e.g. $PATH). That means that the variable substitution might not be perfect.
177        // If there's a default, we need to return the string verbatim as otherwise shellexpand will apply that default for us.
178        if !default.is_empty() {
179            return Ok(Some(format!("${{{var}}}")));
180        }
181        // Else we can just return None and that variable will be left as is.
182        Ok(None)
183    })
184    .ok()?;
185    Some(substituted_string.into_owned())
186}
187
188fn substitute_all_template_variables_in_vec(
189    mut template_strs: Vec<String>,
190    task_variables: &HashMap<String, String>,
191) -> Option<Vec<String>> {
192    for variable in template_strs.iter_mut() {
193        let new_value = substitute_all_template_variables_in_str(&variable, task_variables)?;
194        *variable = new_value;
195    }
196    Some(template_strs)
197}
198
199fn substitute_all_template_variables_in_map(
200    keys_and_values: HashMap<String, String>,
201    task_variables: &HashMap<String, String>,
202) -> Option<HashMap<String, String>> {
203    let mut new_map: HashMap<String, String> = Default::default();
204    for (key, value) in keys_and_values {
205        let new_value = substitute_all_template_variables_in_str(&value, task_variables)?;
206        let new_key = substitute_all_template_variables_in_str(&key, task_variables)?;
207        new_map.insert(new_key, new_value);
208    }
209    Some(new_map)
210}
211
212#[cfg(test)]
213mod tests {
214    use std::{borrow::Cow, path::Path};
215
216    use crate::{TaskVariables, VariableName};
217
218    use super::*;
219
220    const TEST_ID_BASE: &str = "test_base";
221
222    #[test]
223    fn test_resolving_templates_with_blank_command_and_label() {
224        let task_with_all_properties = TaskTemplate {
225            label: "test_label".to_string(),
226            command: "test_command".to_string(),
227            args: vec!["test_arg".to_string()],
228            env: HashMap::from_iter([("test_env_key".to_string(), "test_env_var".to_string())]),
229            ..TaskTemplate::default()
230        };
231
232        for task_with_blank_property in &[
233            TaskTemplate {
234                label: "".to_string(),
235                ..task_with_all_properties.clone()
236            },
237            TaskTemplate {
238                command: "".to_string(),
239                ..task_with_all_properties.clone()
240            },
241            TaskTemplate {
242                label: "".to_string(),
243                command: "".to_string(),
244                ..task_with_all_properties.clone()
245            },
246        ] {
247            assert_eq!(
248                task_with_blank_property.resolve_task(TEST_ID_BASE, TaskContext::default()),
249                None,
250                "should not resolve task with blank label and/or command: {task_with_blank_property:?}"
251            );
252        }
253    }
254
255    #[test]
256    fn test_template_cwd_resolution() {
257        let task_without_cwd = TaskTemplate {
258            cwd: None,
259            label: "test task".to_string(),
260            command: "echo 4".to_string(),
261            ..TaskTemplate::default()
262        };
263
264        let resolved_task = |task_template: &TaskTemplate, task_cx| {
265            let resolved_task = task_template
266                .resolve_task(TEST_ID_BASE, task_cx)
267                .unwrap_or_else(|| panic!("failed to resolve task {task_without_cwd:?}"));
268            resolved_task
269                .resolved
270                .clone()
271                .unwrap_or_else(|| {
272                    panic!("failed to get resolve data for resolved task. Template: {task_without_cwd:?} Resolved: {resolved_task:?}")
273                })
274        };
275
276        assert_eq!(
277            resolved_task(
278                &task_without_cwd,
279                TaskContext {
280                    cwd: None,
281                    task_variables: TaskVariables::default(),
282                }
283            )
284            .cwd,
285            None,
286            "When neither task nor task context have cwd, it should be None"
287        );
288
289        let context_cwd = Path::new("a").join("b").join("c");
290        assert_eq!(
291            resolved_task(
292                &task_without_cwd,
293                TaskContext {
294                    cwd: Some(context_cwd.clone()),
295                    task_variables: TaskVariables::default(),
296                }
297            )
298            .cwd
299            .as_deref(),
300            Some(context_cwd.as_path()),
301            "TaskContext's cwd should be taken on resolve if task's cwd is None"
302        );
303
304        let task_cwd = Path::new("d").join("e").join("f");
305        let mut task_with_cwd = task_without_cwd.clone();
306        task_with_cwd.cwd = Some(task_cwd.display().to_string());
307        let task_with_cwd = task_with_cwd;
308
309        assert_eq!(
310            resolved_task(
311                &task_with_cwd,
312                TaskContext {
313                    cwd: None,
314                    task_variables: TaskVariables::default(),
315                }
316            )
317            .cwd
318            .as_deref(),
319            Some(task_cwd.as_path()),
320            "TaskTemplate's cwd should be taken on resolve if TaskContext's cwd is None"
321        );
322
323        assert_eq!(
324            resolved_task(
325                &task_with_cwd,
326                TaskContext {
327                    cwd: Some(context_cwd.clone()),
328                    task_variables: TaskVariables::default(),
329                }
330            )
331            .cwd
332            .as_deref(),
333            Some(task_cwd.as_path()),
334            "TaskTemplate's cwd should be taken on resolve if TaskContext's cwd is not None"
335        );
336    }
337
338    #[test]
339    fn test_template_variables_resolution() {
340        let custom_variable_1 = VariableName::Custom(Cow::Borrowed("custom_variable_1"));
341        let custom_variable_2 = VariableName::Custom(Cow::Borrowed("custom_variable_2"));
342        let long_value = "01".repeat(MAX_DISPLAY_VARIABLE_LENGTH * 2);
343        let all_variables = [
344            (VariableName::Row, "1234".to_string()),
345            (VariableName::Column, "5678".to_string()),
346            (VariableName::File, "test_file".to_string()),
347            (VariableName::SelectedText, "test_selected_text".to_string()),
348            (VariableName::Symbol, long_value.clone()),
349            (VariableName::WorktreeRoot, "/test_root/".to_string()),
350            (
351                custom_variable_1.clone(),
352                "test_custom_variable_1".to_string(),
353            ),
354            (
355                custom_variable_2.clone(),
356                "test_custom_variable_2".to_string(),
357            ),
358        ];
359
360        let task_with_all_variables = TaskTemplate {
361            label: format!(
362                "test label for {} and {}",
363                VariableName::Row.template_value(),
364                VariableName::Symbol.template_value(),
365            ),
366            command: format!(
367                "echo {} {}",
368                VariableName::File.template_value(),
369                VariableName::Symbol.template_value(),
370            ),
371            args: vec![
372                format!("arg1 {}", VariableName::SelectedText.template_value()),
373                format!("arg2 {}", VariableName::Column.template_value()),
374                format!("arg3 {}", VariableName::Symbol.template_value()),
375            ],
376            env: HashMap::from_iter([
377                ("test_env_key".to_string(), "test_env_var".to_string()),
378                (
379                    "env_key_1".to_string(),
380                    VariableName::WorktreeRoot.template_value(),
381                ),
382                (
383                    "env_key_2".to_string(),
384                    format!(
385                        "env_var_2_{}_{}",
386                        custom_variable_1.template_value(),
387                        custom_variable_2.template_value()
388                    ),
389                ),
390                (
391                    "env_key_3".to_string(),
392                    format!("env_var_3_{}", VariableName::Symbol.template_value()),
393                ),
394            ]),
395            ..TaskTemplate::default()
396        };
397
398        let mut first_resolved_id = None;
399        for i in 0..15 {
400            let resolved_task = task_with_all_variables.resolve_task(
401                TEST_ID_BASE,
402                TaskContext {
403                    cwd: None,
404                    task_variables: TaskVariables::from_iter(all_variables.clone()),
405                },
406            ).unwrap_or_else(|| panic!("Should successfully resolve task {task_with_all_variables:?} with variables {all_variables:?}"));
407
408            match &first_resolved_id {
409                None => first_resolved_id = Some(resolved_task.id),
410                Some(first_id) => assert_eq!(
411                    &resolved_task.id, first_id,
412                    "Step {i}, for the same task template and context, there should be the same resolved task id"
413                ),
414            }
415
416            assert_eq!(
417                resolved_task.original_task, task_with_all_variables,
418                "Resolved task should store its template without changes"
419            );
420            assert_eq!(
421                resolved_task.resolved_label,
422                format!("test label for 1234 and {long_value}"),
423                "Resolved task label should be substituted with variables and those should not be shortened"
424            );
425
426            let spawn_in_terminal = resolved_task
427                .resolved
428                .as_ref()
429                .expect("should have resolved a spawn in terminal task");
430            assert_eq!(
431                spawn_in_terminal.label,
432                format!(
433                    "test label for 1234 and …{}",
434                    &long_value[..=MAX_DISPLAY_VARIABLE_LENGTH]
435                ),
436                "Human-readable label should have long substitutions trimmed"
437            );
438            assert_eq!(
439                spawn_in_terminal.command,
440                format!("echo test_file {long_value}"),
441                "Command should be substituted with variables and those should not be shortened"
442            );
443            assert_eq!(
444                spawn_in_terminal.args,
445                &[
446                    "arg1 test_selected_text",
447                    "arg2 5678",
448                    &format!("arg3 {long_value}")
449                ],
450                "Args should be substituted with variables and those should not be shortened"
451            );
452
453            assert_eq!(
454                spawn_in_terminal
455                    .env
456                    .get("test_env_key")
457                    .map(|s| s.as_str()),
458                Some("test_env_var")
459            );
460            assert_eq!(
461                spawn_in_terminal.env.get("env_key_1").map(|s| s.as_str()),
462                Some("/test_root/")
463            );
464            assert_eq!(
465                spawn_in_terminal.env.get("env_key_2").map(|s| s.as_str()),
466                Some("env_var_2_test_custom_variable_1_test_custom_variable_2")
467            );
468            assert_eq!(
469                spawn_in_terminal.env.get("env_key_3"),
470                Some(&format!("env_var_3_{long_value}")),
471                "Env vars should be substituted with variables and those should not be shortened"
472            );
473        }
474
475        for i in 0..all_variables.len() {
476            let mut not_all_variables = all_variables.to_vec();
477            let removed_variable = not_all_variables.remove(i);
478            let resolved_task_attempt = task_with_all_variables.resolve_task(
479                TEST_ID_BASE,
480                TaskContext {
481                    cwd: None,
482                    task_variables: TaskVariables::from_iter(not_all_variables),
483                },
484            );
485            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})");
486        }
487    }
488
489    #[test]
490    fn test_can_resolve_free_variables() {
491        let task = TaskTemplate {
492            label: "My task".into(),
493            command: "echo".into(),
494            args: vec!["$PATH".into()],
495            ..Default::default()
496        };
497        let resolved = task
498            .resolve_task(TEST_ID_BASE, TaskContext::default())
499            .unwrap()
500            .resolved
501            .unwrap();
502        assert_eq!(resolved.label, task.label);
503        assert_eq!(resolved.command, task.command);
504        assert_eq!(resolved.args, task.args);
505    }
506
507    #[test]
508    fn test_errors_on_missing_zed_variable() {
509        let task = TaskTemplate {
510            label: "My task".into(),
511            command: "echo".into(),
512            args: vec!["$ZED_VARIABLE".into()],
513            ..Default::default()
514        };
515        assert!(task
516            .resolve_task(TEST_ID_BASE, TaskContext::default())
517            .is_none());
518    }
519}