Stricten Zed Task variable API (#10163)

Kirill Bulatov created

Introduce `VariableName` enum to simplify Zed task templating
management: now all the variables can be looked up statically and can be
checked/modified in a centralized way: e.g. `ZED_` prefix is now added
for all such custom vars.

Release Notes:

- N/A

Change summary

crates/language/src/task_context.rs  | 24 +++++----
crates/languages/src/elixir.rs       | 32 ++++++++---
crates/languages/src/rust.rs         | 34 ++++++++----
crates/project/src/task_inventory.rs |  2 
crates/task/src/lib.rs               | 76 +++++++++++++++++++++++++++--
crates/task/src/oneshot_source.rs    |  4 
crates/task/src/static_source.rs     |  7 +-
crates/task/src/vscode_format.rs     | 19 +++++-
crates/tasks_ui/src/lib.rs           | 77 ++++++++++++++---------------
crates/tasks_ui/src/modal.rs         |  4 -
10 files changed, 188 insertions(+), 91 deletions(-)

Detailed changes

crates/language/src/task_context.rs 🔗

@@ -3,13 +3,17 @@ use crate::{LanguageRegistry, Location};
 use anyhow::Result;
 use gpui::{AppContext, Context, Model};
 use std::sync::Arc;
-use task::{static_source::tasks_for, static_source::TaskDefinitions, TaskSource, TaskVariables};
+use task::{
+    static_source::{tasks_for, TaskDefinitions},
+    TaskSource, TaskVariables, VariableName,
+};
 
 /// Language Contexts are used by Zed tasks to extract information about source file.
 pub trait ContextProvider: Send + Sync {
     fn build_context(&self, _: Location, _: &mut AppContext) -> Result<TaskVariables> {
         Ok(TaskVariables::default())
     }
+
     fn associated_tasks(&self) -> Option<TaskDefinitions> {
         None
     }
@@ -29,18 +33,16 @@ impl ContextProvider for SymbolContextProvider {
             .read(cx)
             .snapshot()
             .symbols_containing(location.range.start, None);
-        let symbol = symbols.and_then(|symbols| {
-            symbols.last().map(|symbol| {
-                let range = symbol
-                    .name_ranges
-                    .last()
-                    .cloned()
-                    .unwrap_or(0..symbol.text.len());
-                symbol.text[range].to_string()
-            })
+        let symbol = symbols.unwrap_or_default().last().map(|symbol| {
+            let range = symbol
+                .name_ranges
+                .last()
+                .cloned()
+                .unwrap_or(0..symbol.text.len());
+            symbol.text[range].to_string()
         });
         Ok(TaskVariables::from_iter(
-            symbol.map(|symbol| ("ZED_SYMBOL".to_string(), symbol)),
+            Some(VariableName::Symbol).zip(symbol),
         ))
     }
 }

crates/languages/src/elixir.rs 🔗

@@ -20,7 +20,10 @@ use std::{
         Arc,
     },
 };
-use task::static_source::{Definition, TaskDefinitions};
+use task::{
+    static_source::{Definition, TaskDefinitions},
+    VariableName,
+};
 use util::{
     fs::remove_matching,
     github::{latest_github_release, GitHubLspBinaryVersion},
@@ -557,25 +560,32 @@ pub(super) fn elixir_task_context() -> ContextProviderWithTasks {
             label: "Elixir: test suite".to_owned(),
             command: "mix".to_owned(),
             args: vec!["test".to_owned()],
-            ..Default::default()
+            ..Definition::default()
         },
         Definition {
             label: "Elixir: failed tests suite".to_owned(),
             command: "mix".to_owned(),
             args: vec!["test".to_owned(), "--failed".to_owned()],
-            ..Default::default()
+            ..Definition::default()
         },
         Definition {
             label: "Elixir: test file".to_owned(),
             command: "mix".to_owned(),
-            args: vec!["test".to_owned(), "$ZED_FILE".to_owned()],
-            ..Default::default()
+            args: vec!["test".to_owned(), VariableName::Symbol.template_value()],
+            ..Definition::default()
         },
         Definition {
             label: "Elixir: test at current line".to_owned(),
             command: "mix".to_owned(),
-            args: vec!["test".to_owned(), "$ZED_FILE:$ZED_ROW".to_owned()],
-            ..Default::default()
+            args: vec![
+                "test".to_owned(),
+                format!(
+                    "{}:{}",
+                    VariableName::File.template_value(),
+                    VariableName::Row.template_value()
+                ),
+            ],
+            ..Definition::default()
         },
         Definition {
             label: "Elixir: break line".to_owned(),
@@ -585,9 +595,13 @@ pub(super) fn elixir_task_context() -> ContextProviderWithTasks {
                 "mix".to_owned(),
                 "test".to_owned(),
                 "-b".to_owned(),
-                "$ZED_FILE:$ZED_ROW".to_owned(),
+                format!(
+                    "{}:{}",
+                    VariableName::File.template_value(),
+                    VariableName::Row.template_value()
+                ),
             ],
-            ..Default::default()
+            ..Definition::default()
         },
     ]))
 }

crates/languages/src/rust.rs 🔗

@@ -13,7 +13,7 @@ use smol::fs::{self, File};
 use std::{any::Any, borrow::Cow, env::consts, path::PathBuf, sync::Arc};
 use task::{
     static_source::{Definition, TaskDefinitions},
-    TaskVariables,
+    TaskVariables, VariableName,
 };
 use util::{
     fs::remove_matching,
@@ -322,6 +322,9 @@ impl LspAdapter for RustLspAdapter {
 
 pub(crate) struct RustContextProvider;
 
+const RUST_PACKAGE_TASK_VARIABLE: VariableName =
+    VariableName::Custom(Cow::Borrowed("RUST_PACKAGE"));
+
 impl ContextProvider for RustContextProvider {
     fn build_context(
         &self,
@@ -347,19 +350,24 @@ impl ContextProvider for RustContextProvider {
                 .ok();
 
             if let Some(package_name) = package_name {
-                context.0.insert("ZED_PACKAGE".to_owned(), package_name);
+                context.insert(RUST_PACKAGE_TASK_VARIABLE.clone(), package_name);
             }
         }
 
         Ok(context)
     }
+
     fn associated_tasks(&self) -> Option<TaskDefinitions> {
         Some(TaskDefinitions(vec![
             Definition {
                 label: "Rust: Test current crate".to_owned(),
                 command: "cargo".into(),
-                args: vec!["test".into(), "-p".into(), "$ZED_PACKAGE".into()],
-                ..Default::default()
+                args: vec![
+                    "test".into(),
+                    "-p".into(),
+                    RUST_PACKAGE_TASK_VARIABLE.template_value(),
+                ],
+                ..Definition::default()
             },
             Definition {
                 label: "Rust: Test current function".to_owned(),
@@ -367,29 +375,33 @@ impl ContextProvider for RustContextProvider {
                 args: vec![
                     "test".into(),
                     "-p".into(),
-                    "$ZED_PACKAGE".into(),
+                    RUST_PACKAGE_TASK_VARIABLE.template_value(),
                     "--".into(),
-                    "$ZED_SYMBOL".into(),
+                    VariableName::Symbol.template_value(),
                 ],
-                ..Default::default()
+                ..Definition::default()
             },
             Definition {
                 label: "Rust: cargo run".into(),
                 command: "cargo".into(),
                 args: vec!["run".into()],
-                ..Default::default()
+                ..Definition::default()
             },
             Definition {
                 label: "Rust: cargo check current crate".into(),
                 command: "cargo".into(),
-                args: vec!["check".into(), "-p".into(), "$ZED_PACKAGE".into()],
-                ..Default::default()
+                args: vec![
+                    "check".into(),
+                    "-p".into(),
+                    RUST_PACKAGE_TASK_VARIABLE.template_value(),
+                ],
+                ..Definition::default()
             },
             Definition {
                 label: "Rust: cargo check workspace".into(),
                 command: "cargo".into(),
                 args: vec!["check".into(), "--workspace".into()],
-                ..Default::default()
+                ..Definition::default()
             },
         ]))
     }

crates/project/src/task_inventory.rs 🔗

@@ -258,7 +258,7 @@ pub mod test_inventory {
             None
         }
 
-        fn exec(&self, _cwd: TaskContext) -> Option<task::SpawnInTerminal> {
+        fn prepare_exec(&self, _cwd: TaskContext) -> Option<task::SpawnInTerminal> {
             None
         }
     }

crates/task/src/lib.rs 🔗

@@ -9,6 +9,7 @@ use collections::HashMap;
 use gpui::ModelContext;
 use static_source::RevealStrategy;
 use std::any::Any;
+use std::borrow::Cow;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 pub use vscode_format::VsCodeTaskFile;
@@ -41,15 +42,78 @@ pub struct SpawnInTerminal {
     pub reveal: RevealStrategy,
 }
 
-type VariableName = String;
-type VariableValue = String;
+/// Variables, available for use in [`TaskContext`] when a Zed's task gets turned into real command.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub enum VariableName {
+    /// An absolute path of the currently opened file.
+    File,
+    /// An absolute path of the currently opened worktree, that contains the file.
+    WorktreeRoot,
+    /// A symbol text, that contains latest cursor/selection position.
+    Symbol,
+    /// A row with the latest cursor/selection position.
+    Row,
+    /// A column with the latest cursor/selection position.
+    Column,
+    /// Text from the latest selection.
+    SelectedText,
+    /// Custom variable, provided by the plugin or other external source.
+    /// Will be printed with `ZED_` prefix to avoid potential conflicts with other variables.
+    Custom(Cow<'static, str>),
+}
+
+impl VariableName {
+    /// Generates a `$VARIABLE`-like string value to be used in templates.
+    /// Custom variables are wrapped in `${}` to avoid substitution issues with whitespaces.
+    pub fn template_value(&self) -> String {
+        if matches!(self, Self::Custom(_)) {
+            format!("${{{self}}}")
+        } else {
+            format!("${self}")
+        }
+    }
+}
+
+impl std::fmt::Display for VariableName {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match self {
+            Self::File => write!(f, "ZED_FILE"),
+            Self::WorktreeRoot => write!(f, "ZED_WORKTREE_ROOT"),
+            Self::Symbol => write!(f, "ZED_SYMBOL"),
+            Self::Row => write!(f, "ZED_ROW"),
+            Self::Column => write!(f, "ZED_COLUMN"),
+            Self::SelectedText => write!(f, "ZED_SELECTED_TEXT"),
+            Self::Custom(s) => write!(f, "ZED_{s}"),
+        }
+    }
+}
 
 /// Container for predefined environment variables that describe state of Zed at the time the task was spawned.
 #[derive(Clone, Debug, Default, PartialEq, Eq)]
-pub struct TaskVariables(pub HashMap<VariableName, VariableValue>);
+pub struct TaskVariables(HashMap<VariableName, String>);
+
+impl TaskVariables {
+    /// Converts the container into a map of environment variables and their values.
+    fn into_env_variables(self) -> HashMap<String, String> {
+        self.0
+            .into_iter()
+            .map(|(name, value)| (name.to_string(), value))
+            .collect()
+    }
+
+    /// Inserts another variable into the container, overwriting the existing one if it already exists — in this case, the old value is returned.
+    pub fn insert(&mut self, variable: VariableName, value: String) -> Option<String> {
+        self.0.insert(variable, value)
+    }
+
+    /// Extends the container with another one, overwriting the existing variables on collision.
+    pub fn extend(&mut self, other: Self) {
+        self.0.extend(other.0);
+    }
+}
 
-impl FromIterator<(String, String)> for TaskVariables {
-    fn from_iter<T: IntoIterator<Item = (String, String)>>(iter: T) -> Self {
+impl FromIterator<(VariableName, String)> for TaskVariables {
+    fn from_iter<T: IntoIterator<Item = (VariableName, String)>>(iter: T) -> Self {
         Self(HashMap::from_iter(iter))
     }
 }
@@ -74,7 +138,7 @@ pub trait Task {
     fn cwd(&self) -> Option<&str>;
     /// Sets up everything needed to spawn the task in the given directory (`cwd`).
     /// If a task is intended to be spawned in the terminal, it should return the corresponding struct filled with the data necessary.
-    fn exec(&self, cx: TaskContext) -> Option<SpawnInTerminal>;
+    fn prepare_exec(&self, cx: TaskContext) -> Option<SpawnInTerminal>;
 }
 
 /// [`Source`] produces tasks that can be scheduled.

crates/task/src/oneshot_source.rs 🔗

@@ -36,7 +36,7 @@ impl Task for OneshotTask {
         None
     }
 
-    fn exec(&self, cx: TaskContext) -> Option<SpawnInTerminal> {
+    fn prepare_exec(&self, cx: TaskContext) -> Option<SpawnInTerminal> {
         if self.id().0.is_empty() {
             return None;
         }
@@ -50,7 +50,7 @@ impl Task for OneshotTask {
             command: self.id().0.clone(),
             args: vec![],
             cwd,
-            env: task_variables.0,
+            env: task_variables.into_env_variables(),
             use_new_terminal: Default::default(),
             allow_concurrent_runs: Default::default(),
             reveal: RevealStrategy::default(),

crates/task/src/static_source.rs 🔗

@@ -42,23 +42,24 @@ pub fn tasks_for(tasks: TaskDefinitions, id_base: &str) -> Vec<Arc<dyn Task>> {
 }
 
 impl Task for StaticTask {
-    fn exec(&self, cx: TaskContext) -> Option<SpawnInTerminal> {
+    fn prepare_exec(&self, cx: TaskContext) -> Option<SpawnInTerminal> {
         let TaskContext {
             cwd,
             task_variables,
         } = cx;
+        let task_variables = task_variables.into_env_variables();
         let cwd = self
             .definition
             .cwd
             .clone()
             .and_then(|path| {
-                subst::substitute(&path, &task_variables.0)
+                subst::substitute(&path, &task_variables)
                     .map(Into::into)
                     .ok()
             })
             .or(cwd);
         let mut definition_env = self.definition.env.clone();
-        definition_env.extend(task_variables.0);
+        definition_env.extend(task_variables);
         Some(SpawnInTerminal {
             id: self.id.clone(),
             cwd,

crates/task/src/vscode_format.rs 🔗

@@ -3,7 +3,10 @@ use collections::HashMap;
 use serde::Deserialize;
 use util::ResultExt;
 
-use crate::static_source::{Definition, TaskDefinitions};
+use crate::{
+    static_source::{Definition, TaskDefinitions},
+    VariableName,
+};
 
 #[derive(Clone, Debug, Deserialize, PartialEq)]
 #[serde(rename_all = "camelCase")]
@@ -129,10 +132,16 @@ impl TryFrom<VsCodeTaskFile> for TaskDefinitions {
 
     fn try_from(value: VsCodeTaskFile) -> Result<Self, Self::Error> {
         let replacer = EnvVariableReplacer::new(HashMap::from_iter([
-            ("workspaceFolder".to_owned(), "ZED_WORKTREE_ROOT".to_owned()),
-            ("file".to_owned(), "ZED_FILE".to_owned()),
-            ("lineNumber".to_owned(), "ZED_ROW".to_owned()),
-            ("selectedText".to_owned(), "ZED_SELECTED_TEXT".to_owned()),
+            (
+                "workspaceFolder".to_owned(),
+                VariableName::WorktreeRoot.to_string(),
+            ),
+            ("file".to_owned(), VariableName::File.to_string()),
+            ("lineNumber".to_owned(), VariableName::Row.to_string()),
+            (
+                "selectedText".to_owned(),
+                VariableName::SelectedText.to_string(),
+            ),
         ]));
         let definitions = value
             .tasks

crates/tasks_ui/src/lib.rs 🔗

@@ -5,7 +5,7 @@ use gpui::{AppContext, ViewContext, WindowContext};
 use language::Point;
 use modal::{Spawn, TasksModal};
 use project::{Location, WorktreeId};
-use task::{Task, TaskContext, TaskVariables};
+use task::{Task, TaskContext, TaskVariables, VariableName};
 use util::ResultExt;
 use workspace::Workspace;
 
@@ -30,7 +30,6 @@ pub fn init(cx: &mut AppContext) {
                         } else {
                             old_context
                         };
-
                         schedule_task(workspace, task.as_ref(), task_context, false, cx)
                     };
                 });
@@ -40,17 +39,17 @@ pub fn init(cx: &mut AppContext) {
 }
 
 fn spawn_task_or_modal(workspace: &mut Workspace, action: &Spawn, cx: &mut ViewContext<Workspace>) {
-    let inventory = workspace.project().read(cx).task_inventory().clone();
-    let workspace_handle = workspace.weak_handle();
-    let cwd = task_cwd(workspace, cx).log_err().flatten();
-    let task_context = task_context(workspace, cwd, cx);
-    if let Some(name) = action.task_name.clone() {
-        // Do not actually show the modal.
-        spawn_task_with_name(name.clone(), cx);
-    } else {
-        workspace.toggle_modal(cx, |cx| {
-            TasksModal::new(inventory, task_context, workspace_handle, cx)
-        })
+    match &action.task_name {
+        Some(name) => spawn_task_with_name(name.clone(), cx),
+        None => {
+            let inventory = workspace.project().read(cx).task_inventory().clone();
+            let workspace_handle = workspace.weak_handle();
+            let cwd = task_cwd(workspace, cx).log_err().flatten();
+            let task_context = task_context(workspace, cwd, cx);
+            workspace.toggle_modal(cx, |cx| {
+                TasksModal::new(inventory, task_context, workspace_handle, cx)
+            })
+        }
     }
 }
 
@@ -157,20 +156,18 @@ fn task_context(
                 let selected_text = buffer.read(cx).chars_for_range(selection_range).collect();
 
                 let mut task_variables = TaskVariables::from_iter([
-                    ("ZED_ROW".into(), row.to_string()),
-                    ("ZED_COLUMN".into(), column.to_string()),
-                    ("ZED_SELECTED_TEXT".into(), selected_text),
+                    (VariableName::Row, row.to_string()),
+                    (VariableName::Column, column.to_string()),
+                    (VariableName::SelectedText, selected_text),
                 ]);
                 if let Some(path) = current_file {
-                    task_variables.0.insert("ZED_FILE".into(), path);
+                    task_variables.insert(VariableName::File, path);
                 }
                 if let Some(worktree_path) = worktree_path {
-                    task_variables
-                        .0
-                        .insert("ZED_WORKTREE_ROOT".into(), worktree_path);
+                    task_variables.insert(VariableName::WorktreeRoot, worktree_path);
                 }
                 if let Some(language_context) = context {
-                    task_variables.0.extend(language_context.0);
+                    task_variables.extend(language_context);
                 }
 
                 Some(TaskContext {
@@ -198,7 +195,7 @@ fn schedule_task(
     omit_history: bool,
     cx: &mut ViewContext<'_, Workspace>,
 ) {
-    let spawn_in_terminal = task.exec(task_cx.clone());
+    let spawn_in_terminal = task.prepare_exec(task_cx.clone());
     if let Some(spawn_in_terminal) = spawn_in_terminal {
         if !omit_history {
             workspace.project().update(cx, |project, cx| {
@@ -255,7 +252,7 @@ mod tests {
     use language::{Language, LanguageConfig, SymbolContextProvider};
     use project::{FakeFs, Project, TaskSourceKind};
     use serde_json::json;
-    use task::{oneshot_source::OneshotSource, TaskContext, TaskVariables};
+    use task::{oneshot_source::OneshotSource, TaskContext, TaskVariables, VariableName};
     use ui::VisualContext;
     use workspace::{AppState, Workspace};
 
@@ -363,11 +360,11 @@ mod tests {
                 TaskContext {
                     cwd: Some("/dir".into()),
                     task_variables: TaskVariables::from_iter([
-                        ("ZED_FILE".into(), "/dir/rust/b.rs".into()),
-                        ("ZED_WORKTREE_ROOT".into(), "/dir".into()),
-                        ("ZED_ROW".into(), "1".into()),
-                        ("ZED_COLUMN".into(), "1".into()),
-                        ("ZED_SELECTED_TEXT".into(), "".into())
+                        (VariableName::File, "/dir/rust/b.rs".into()),
+                        (VariableName::WorktreeRoot, "/dir".into()),
+                        (VariableName::Row, "1".into()),
+                        (VariableName::Column, "1".into()),
+                        (VariableName::SelectedText, "".into())
                     ])
                 }
             );
@@ -380,12 +377,12 @@ mod tests {
                 TaskContext {
                     cwd: Some("/dir".into()),
                     task_variables: TaskVariables::from_iter([
-                        ("ZED_FILE".into(), "/dir/rust/b.rs".into()),
-                        ("ZED_WORKTREE_ROOT".into(), "/dir".into()),
-                        ("ZED_SYMBOL".into(), "this_is_a_rust_file".into()),
-                        ("ZED_ROW".into(), "1".into()),
-                        ("ZED_COLUMN".into(), "15".into()),
-                        ("ZED_SELECTED_TEXT".into(), "is_i".into()),
+                        (VariableName::File, "/dir/rust/b.rs".into()),
+                        (VariableName::WorktreeRoot, "/dir".into()),
+                        (VariableName::Row, "1".into()),
+                        (VariableName::Column, "15".into()),
+                        (VariableName::SelectedText, "is_i".into()),
+                        (VariableName::Symbol, "this_is_a_rust_file".into()),
                     ])
                 }
             );
@@ -397,12 +394,12 @@ mod tests {
                 TaskContext {
                     cwd: Some("/dir".into()),
                     task_variables: TaskVariables::from_iter([
-                        ("ZED_FILE".into(), "/dir/a.ts".into()),
-                        ("ZED_WORKTREE_ROOT".into(), "/dir".into()),
-                        ("ZED_SYMBOL".into(), "this_is_a_test".into()),
-                        ("ZED_ROW".into(), "1".into()),
-                        ("ZED_COLUMN".into(), "1".into()),
-                        ("ZED_SELECTED_TEXT".into(), "".into()),
+                        (VariableName::File, "/dir/a.ts".into()),
+                        (VariableName::WorktreeRoot, "/dir".into()),
+                        (VariableName::Row, "1".into()),
+                        (VariableName::Column, "1".into()),
+                        (VariableName::SelectedText, "".into()),
+                        (VariableName::Symbol, "this_is_a_test".into()),
                     ])
                 }
             );

crates/tasks_ui/src/modal.rs 🔗

@@ -366,9 +366,7 @@ impl PickerDelegate for TasksModalDelegate {
         let task_index = self.matches.get(self.selected_index())?.candidate_id;
         let tasks = self.candidates.as_ref()?;
         let (_, task) = tasks.get(task_index)?;
-        // .exec doesn't actually spawn anything; it merely prepares a spawning command,
-        // which we can use for substitution.
-        let mut spawn_prompt = task.exec(self.task_context.clone())?;
+        let mut spawn_prompt = task.prepare_exec(self.task_context.clone())?;
         if !spawn_prompt.args.is_empty() {
             spawn_prompt.command.push(' ');
             spawn_prompt