tasks: Expose captured variables to ContextProvider (#12134)

This PR changes the interface of ContextProvider, allowing it to inspect
*all* variables set so far during the process of building
`TaskVariables`. This makes it possible to capture e.g. an identifier in
tree-sitter query, process it and then export it as a task variable.

Notably, the list of variables includes captures prefixed with leading
underscore; they are removed after all calls to `build_context`, but it
makes it possible to capture something and then conditionally preserve
it (and perhaps modify it).

Release Notes:

- N/A

crates/editor/src/editor.rs

@@ -406,6 +406,7 @@ struct RunnableTasks {
     templates: Vec<(TaskSourceKind, TaskTemplate)>,
     // We need the column at which the task context evaluation should take place.
     column: u32,
+    // Values of all named captures, including those starting with '_'
     extra_variables: HashMap<String, String>,
 }
 
@@ -3973,7 +3974,7 @@ impl Editor {
 
                     this.completion_tasks.clear();
                     this.discard_inline_completion(false, cx);
-                    let task_context = tasks.as_ref().zip(this.workspace.clone()).and_then(
+                    let tasks = tasks.as_ref().zip(this.workspace.clone()).and_then(
                         |(tasks, (workspace, _))| {
                             let position = Point::new(buffer_row, tasks.1.column);
                             let range_start = buffer.read(cx).anchor_at(position, Bias::Right);
@@ -3981,43 +3982,45 @@ impl Editor {
                                 buffer: buffer.clone(),
                                 range: range_start..range_start,
                             };
+                            // Fill in the environmental variables from the tree-sitter captures
+                            let mut captured_task_variables = TaskVariables::default();
+                            for (capture_name, value) in tasks.1.extra_variables.clone() {
+                                captured_task_variables.insert(
+                                    task::VariableName::Custom(capture_name.into()),
+                                    value.clone(),
+                                );
+                            }
+
                             workspace
                                 .update(cx, |workspace, cx| {
-                                    tasks::task_context_for_location(workspace, location, cx)
+                                    tasks::task_context_for_location(
+                                        captured_task_variables,
+                                        workspace,
+                                        location,
+                                        cx,
+                                    )
                                 })
                                 .ok()
                                 .flatten()
+                                .map(|task_context| {
+                                    Arc::new(ResolvedTasks {
+                                        templates: tasks
+                                            .1
+                                            .templates
+                                            .iter()
+                                            .filter_map(|(kind, template)| {
+                                                template
+                                                    .resolve_task(&kind.to_id_base(), &task_context)
+                                                    .map(|task| (kind.clone(), task))
+                                            })
+                                            .collect(),
+                                        position: snapshot.buffer_snapshot.anchor_before(
+                                            Point::new(multibuffer_point.row, tasks.1.column),
+                                        ),
+                                    })
+                                })
                         },
                     );
-                    let tasks = tasks.zip(task_context).map(|(tasks, mut task_context)| {
-                        // Fill in the environmental variables from the tree-sitter captures
-                        let mut additional_task_variables = TaskVariables::default();
-                        for (capture_name, value) in tasks.1.extra_variables.clone() {
-                            additional_task_variables.insert(
-                                task::VariableName::Custom(capture_name.into()),
-                                value.clone(),
-                            );
-                        }
-                        task_context
-                            .task_variables
-                            .extend(additional_task_variables);
-
-                        Arc::new(ResolvedTasks {
-                            templates: tasks
-                                .1
-                                .templates
-                                .iter()
-                                .filter_map(|(kind, template)| {
-                                    template
-                                        .resolve_task(&kind.to_id_base(), &task_context)
-                                        .map(|task| (kind.clone(), task))
-                                })
-                                .collect(),
-                            position: snapshot
-                                .buffer_snapshot
-                                .anchor_before(Point::new(multibuffer_point.row, tasks.1.column)),
-                        })
-                    });
                     let spawn_straight_away = tasks
                         .as_ref()
                         .map_or(false, |tasks| tasks.templates.len() == 1)

crates/editor/src/tasks.rs

@@ -1,17 +1,16 @@
 use crate::Editor;
 
-use std::{path::Path, sync::Arc};
-
 use anyhow::Context;
-use gpui::WindowContext;
+use gpui::{Model, WindowContext};
-use language::{BasicContextProvider, ContextProvider};
+use language::ContextProvider;
-use project::{Location, WorktreeId};
+use project::{BasicContextProvider, Location, Project};
 use task::{TaskContext, TaskVariables, VariableName};
 use text::Point;
 use util::ResultExt;
 use workspace::Workspace;
 
 pub(crate) fn task_context_for_location(
+    captured_variables: TaskVariables,
     workspace: &Workspace,
     location: Location,
     cx: &mut WindowContext<'_>,
@@ -20,31 +19,16 @@ pub(crate) fn task_context_for_location(
         .log_err()
         .flatten();
 
-    let buffer = location.buffer.clone();
+    let mut task_variables = combine_task_variables(
-    let language_context_provider = buffer
+        captured_variables,
-        .read(cx)
-        .language()
-        .and_then(|language| language.context_provider())
-        .unwrap_or_else(|| Arc::new(BasicContextProvider));
-
-    let worktree_abs_path = buffer
-        .read(cx)
-        .file()
-        .map(|file| WorktreeId::from_usize(file.worktree_id()))
-        .and_then(|worktree_id| {
-            workspace
-                .project()
-                .read(cx)
-                .worktree_for_id(worktree_id, cx)
-                .map(|worktree| worktree.read(cx).abs_path())
-        });
-    let task_variables = combine_task_variables(
-        worktree_abs_path.as_deref(),
         location,
-        language_context_provider.as_ref(),
+        workspace.project().clone(),
         cx,
     )
     .log_err()?;
+    // Remove all custom entries starting with _, as they're not intended for use by the end user.
+    task_variables.sweep();
+
     Some(TaskContext {
         cwd,
         task_variables,
@@ -84,21 +68,21 @@ fn task_context_with_editor(
         buffer,
         range: start..end,
     };
-    task_context_for_location(workspace, location.clone(), cx).map(|mut task_context| {
+    let captured_variables = {
+        let mut variables = TaskVariables::default();
         for range in location
             .buffer
             .read(cx)
             .snapshot()
-            .runnable_ranges(location.range)
+            .runnable_ranges(location.range.clone())
         {
             for (capture_name, value) in range.extra_captures {
-                task_context
+                variables.insert(VariableName::Custom(capture_name.into()), value);
-                    .task_variables
-                    .insert(VariableName::Custom(capture_name.into()), value);
             }
         }
-        task_context
+        variables
-    })
+    };
+    task_context_for_location(captured_variables, workspace, location.clone(), cx)
 }
 
 pub fn task_context(workspace: &Workspace, cx: &mut WindowContext<'_>) -> TaskContext {
@@ -114,24 +98,26 @@ pub fn task_context(workspace: &Workspace, cx: &mut WindowContext<'_>) -> TaskCo
 }
 
 fn combine_task_variables(
-    worktree_abs_path: Option<&Path>,
+    mut captured_variables: TaskVariables,
     location: Location,
-    context_provider: &dyn ContextProvider,
+    project: Model<Project>,
     cx: &mut WindowContext<'_>,
 ) -> anyhow::Result<TaskVariables> {
-    if context_provider.is_basic() {
+    let language_context_provider = location
-        context_provider
+        .buffer
-            .build_context(worktree_abs_path, &location, cx)
+        .read(cx)
-            .context("building basic provider context")
+        .language()
-    } else {
+        .and_then(|language| language.context_provider());
-        let mut basic_context = BasicContextProvider
+    let baseline = BasicContextProvider::new(project)
-            .build_context(worktree_abs_path, &location, cx)
+        .build_context(&captured_variables, &location, cx)
-            .context("building basic default context")?;
+        .context("building basic default context")?;
-        basic_context.extend(
+    captured_variables.extend(baseline);
-            context_provider
+    if let Some(provider) = language_context_provider {
-                .build_context(worktree_abs_path, &location, cx)
+        captured_variables.extend(
+            provider
+                .build_context(&captured_variables, &location, cx)
                 .context("building provider context ")?,
         );
-        Ok(basic_context)
     }
+    Ok(captured_variables)
 }

crates/extension/src/extension_store.rs

@@ -30,10 +30,10 @@ use gpui::{
 };
 use http::{AsyncBody, HttpClient, HttpClientWithUrl};
 use language::{
-    ContextProviderWithTasks, LanguageConfig, LanguageMatcher, LanguageQueries, LanguageRegistry,
+    LanguageConfig, LanguageMatcher, LanguageQueries, LanguageRegistry, QUERY_FILENAME_PREFIXES,
-    QUERY_FILENAME_PREFIXES,
 };
 use node_runtime::NodeRuntime;
+use project::ContextProviderWithTasks;
 use semantic_version::SemanticVersion;
 use serde::{Deserialize, Serialize};
 use settings::Settings;

crates/language/src/buffer.rs

@@ -3024,6 +3024,7 @@ impl BufferSnapshot {
                 tags.sort_by_key(|(range, _)| range == &maximum_range);
                 let split_point = tags.partition_point(|(range, _)| range != &maximum_range);
                 let (extra_captures, tags) = tags.split_at(split_point);
+
                 let extra_captures = extra_captures
                     .into_iter()
                     .map(|(range, name)| {

crates/language/src/language.rs

@@ -58,9 +58,7 @@ use std::{
 };
 use syntax_map::{QueryCursorHandle, SyntaxSnapshot};
 use task::RunnableTag;
-pub use task_context::{
+pub use task_context::{ContextProvider, RunnableRange};
-    BasicContextProvider, ContextProvider, ContextProviderWithTasks, RunnableRange,
-};
 use theme::SyntaxTheme;
 use tree_sitter::{self, wasmtime, Query, QueryCursor, WasmStore};
 
@@ -1016,7 +1014,7 @@ impl Language {
         for (ix, name) in query.capture_names().iter().enumerate() {
             if *name == "run" {
                 run_capture_index = Some(ix as u32);
-            } else if !name.starts_with('_') {
+            } else {
                 runnable_tags.insert(ix as u32, RunnableTag(name.to_string().into()));
             }
         }

crates/language/src/task_context.rs

@@ -1,12 +1,12 @@
-use std::{ops::Range, path::Path};
+use std::ops::Range;
 
 use crate::{Location, Runnable};
 
 use anyhow::Result;
 use collections::HashMap;
 use gpui::AppContext;
-use task::{TaskTemplates, TaskVariables, VariableName};
+use task::{TaskTemplates, TaskVariables};
-use text::{BufferId, Point, ToPoint};
+use text::BufferId;
 
 pub struct RunnableRange {
     pub buffer_id: BufferId,
@@ -22,7 +22,7 @@ pub trait ContextProvider: Send + Sync {
     /// Builds a specific context to be placed on top of the basic one (replacing all conflicting entries) and to be used for task resolving later.
     fn build_context(
         &self,
-        _worktree_abs_path: Option<&Path>,
+        _variables: &TaskVariables,
         _location: &Location,
         _cx: &mut AppContext,
     ) -> Result<TaskVariables> {
@@ -33,100 +33,4 @@ pub trait ContextProvider: Send + Sync {
     fn associated_tasks(&self) -> Option<TaskTemplates> {
         None
     }
-
-    // Determines whether the [`BasicContextProvider`] variables should be filled too (if `false`), or omitted (if `true`).
-    fn is_basic(&self) -> bool {
-        false
-    }
-}
-
-/// A context provided that tries to provide values for all non-custom [`VariableName`] variants for a currently opened file.
-/// Applied as a base for every custom [`ContextProvider`] unless explicitly oped out.
-pub struct BasicContextProvider;
-
-impl ContextProvider for BasicContextProvider {
-    fn is_basic(&self) -> bool {
-        true
-    }
-
-    fn build_context(
-        &self,
-        worktree_abs_path: Option<&Path>,
-        location: &Location,
-        cx: &mut AppContext,
-    ) -> Result<TaskVariables> {
-        let buffer = location.buffer.read(cx);
-        let buffer_snapshot = buffer.snapshot();
-        let symbols = buffer_snapshot.symbols_containing(location.range.start, None);
-        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()
-        });
-
-        let current_file = buffer
-            .file()
-            .and_then(|file| file.as_local())
-            .map(|file| file.abs_path(cx).to_string_lossy().to_string());
-        let Point { row, column } = location.range.start.to_point(&buffer_snapshot);
-        let row = row + 1;
-        let column = column + 1;
-        let selected_text = buffer
-            .chars_for_range(location.range.clone())
-            .collect::<String>();
-
-        let mut task_variables = TaskVariables::from_iter([
-            (VariableName::Row, row.to_string()),
-            (VariableName::Column, column.to_string()),
-        ]);
-
-        if let Some(symbol) = symbol {
-            task_variables.insert(VariableName::Symbol, symbol);
-        }
-        if !selected_text.trim().is_empty() {
-            task_variables.insert(VariableName::SelectedText, selected_text);
-        }
-        if let Some(path) = current_file {
-            task_variables.insert(VariableName::File, path);
-        }
-        if let Some(worktree_path) = worktree_abs_path {
-            task_variables.insert(
-                VariableName::WorktreeRoot,
-                worktree_path.to_string_lossy().to_string(),
-            );
-        }
-
-        Ok(task_variables)
-    }
-}
-
-/// A ContextProvider that doesn't provide any task variables on it's own, though it has some associated tasks.
-pub struct ContextProviderWithTasks {
-    templates: TaskTemplates,
-}
-
-impl ContextProviderWithTasks {
-    pub fn new(definitions: TaskTemplates) -> Self {
-        Self {
-            templates: definitions,
-        }
-    }
-}
-
-impl ContextProvider for ContextProviderWithTasks {
-    fn associated_tasks(&self) -> Option<TaskTemplates> {
-        Some(self.templates.clone())
-    }
-
-    fn build_context(
-        &self,
-        worktree_abs_path: Option<&Path>,
-        location: &Location,
-        cx: &mut AppContext,
-    ) -> Result<TaskVariables> {
-        BasicContextProvider.build_context(worktree_abs_path, location, cx)
-    }
 }

crates/languages/src/bash.rs

@@ -1,4 +1,4 @@
-use language::ContextProviderWithTasks;
+use project::ContextProviderWithTasks;
 use task::{TaskTemplate, TaskTemplates, VariableName};
 
 pub(super) fn bash_task_context() -> ContextProviderWithTasks {

crates/languages/src/go.rs

@@ -16,7 +16,7 @@ use std::{
     borrow::Cow,
     ffi::{OsStr, OsString},
     ops::Range,
-    path::{Path, PathBuf},
+    path::PathBuf,
     str,
     sync::{
         atomic::{AtomicBool, Ordering::SeqCst},
@@ -447,7 +447,7 @@ const GO_PACKAGE_TASK_VARIABLE: VariableName = VariableName::Custom(Cow::Borrowe
 impl ContextProvider for GoContextProvider {
     fn build_context(
         &self,
-        worktree_abs_path: Option<&Path>,
+        variables: &TaskVariables,
         location: &Location,
         cx: &mut gpui::AppContext,
     ) -> Result<TaskVariables> {
@@ -465,7 +465,8 @@ impl ContextProvider for GoContextProvider {
                 // Prefer the relative form `./my-nested-package/is-here` over
                 // absolute path, because it's more readable in the modal, but
                 // the absolute path also works.
-                let package_name = worktree_abs_path
+                let package_name = variables
+                    .get(&VariableName::WorktreeRoot)
                     .and_then(|worktree_abs_path| buffer_dir.strip_prefix(worktree_abs_path).ok())
                     .map(|relative_pkg_dir| {
                         if relative_pkg_dir.as_os_str().is_empty() {

crates/languages/src/python.rs

@@ -1,8 +1,9 @@
 use anyhow::Result;
 use async_trait::async_trait;
-use language::{ContextProviderWithTasks, LanguageServerName, LspAdapter, LspAdapterDelegate};
+use language::{LanguageServerName, LspAdapter, LspAdapterDelegate};
 use lsp::LanguageServerBinary;
 use node_runtime::NodeRuntime;
+use project::ContextProviderWithTasks;
 use std::{
     any::Any,
     ffi::OsString,

crates/languages/src/rust.rs

@@ -328,7 +328,7 @@ const RUST_PACKAGE_TASK_VARIABLE: VariableName =
 impl ContextProvider for RustContextProvider {
     fn build_context(
         &self,
-        _: Option<&Path>,
+        _: &TaskVariables,
         location: &Location,
         cx: &mut gpui::AppContext,
     ) -> Result<TaskVariables> {

crates/project/src/project.rs

@@ -113,7 +113,9 @@ pub use fs::*;
 pub use language::Location;
 #[cfg(any(test, feature = "test-support"))]
 pub use prettier::FORMAT_SUFFIX as TEST_PRETTIER_FORMAT_SUFFIX;
-pub use task_inventory::{Inventory, TaskSourceKind};
+pub use task_inventory::{
+    BasicContextProvider, ContextProviderWithTasks, Inventory, TaskSourceKind,
+};
 pub use worktree::{
     DiagnosticSummary, Entry, EntryKind, File, LocalWorktree, PathChange, ProjectEntryId,
     RepositoryEntry, UpdatedEntriesSet, UpdatedGitRepositoriesSet, Worktree, WorktreeId,

crates/project/src/task_inventory.rs

@@ -6,6 +6,7 @@ use std::{
     sync::Arc,
 };
 
+use anyhow::Result;
 use collections::{btree_map, BTreeMap, VecDeque};
 use futures::{
     channel::mpsc::{unbounded, UnboundedSender},
@@ -13,13 +14,17 @@ use futures::{
 };
 use gpui::{AppContext, Context, Model, ModelContext, Task};
 use itertools::Itertools;
-use language::Language;
+use language::{ContextProvider, Language, Location};
 use task::{
-    static_source::StaticSource, ResolvedTask, TaskContext, TaskId, TaskTemplate, VariableName,
+    static_source::StaticSource, ResolvedTask, TaskContext, TaskId, TaskTemplate, TaskTemplates,
+    TaskVariables, VariableName,
 };
+use text::{Point, ToPoint};
 use util::{post_inc, NumericPrefixWithSuffix};
 use worktree::WorktreeId;
 
+use crate::Project;
+
 /// Inventory tracks available tasks for a given project.
 pub struct Inventory {
     sources: Vec<SourceInInventory>,
@@ -491,6 +496,102 @@ mod test_inventory {
     }
 }
 
+/// A context provided that tries to provide values for all non-custom [`VariableName`] variants for a currently opened file.
+/// Applied as a base for every custom [`ContextProvider`] unless explicitly oped out.
+pub struct BasicContextProvider {
+    project: Model<Project>,
+}
+
+impl BasicContextProvider {
+    pub fn new(project: Model<Project>) -> Self {
+        Self { project }
+    }
+}
+
+impl ContextProvider for BasicContextProvider {
+    fn build_context(
+        &self,
+        _: &TaskVariables,
+        location: &Location,
+        cx: &mut AppContext,
+    ) -> Result<TaskVariables> {
+        let buffer = location.buffer.read(cx);
+        let buffer_snapshot = buffer.snapshot();
+        let symbols = buffer_snapshot.symbols_containing(location.range.start, None);
+        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()
+        });
+
+        let current_file = buffer
+            .file()
+            .and_then(|file| file.as_local())
+            .map(|file| file.abs_path(cx).to_string_lossy().to_string());
+        let Point { row, column } = location.range.start.to_point(&buffer_snapshot);
+        let row = row + 1;
+        let column = column + 1;
+        let selected_text = buffer
+            .chars_for_range(location.range.clone())
+            .collect::<String>();
+
+        let mut task_variables = TaskVariables::from_iter([
+            (VariableName::Row, row.to_string()),
+            (VariableName::Column, column.to_string()),
+        ]);
+
+        if let Some(symbol) = symbol {
+            task_variables.insert(VariableName::Symbol, symbol);
+        }
+        if !selected_text.trim().is_empty() {
+            task_variables.insert(VariableName::SelectedText, selected_text);
+        }
+        if let Some(path) = current_file {
+            task_variables.insert(VariableName::File, path);
+        }
+
+        let worktree_abs_path = buffer
+            .file()
+            .map(|file| WorktreeId::from_usize(file.worktree_id()))
+            .and_then(|worktree_id| {
+                self.project
+                    .read(cx)
+                    .worktree_for_id(worktree_id, cx)
+                    .map(|worktree| worktree.read(cx).abs_path())
+            });
+        if let Some(worktree_path) = worktree_abs_path {
+            task_variables.insert(
+                VariableName::WorktreeRoot,
+                worktree_path.to_string_lossy().to_string(),
+            );
+        }
+
+        Ok(task_variables)
+    }
+}
+
+/// A ContextProvider that doesn't provide any task variables on it's own, though it has some associated tasks.
+pub struct ContextProviderWithTasks {
+    templates: TaskTemplates,
+}
+
+impl ContextProviderWithTasks {
+    pub fn new(definitions: TaskTemplates) -> Self {
+        Self {
+            templates: definitions,
+        }
+    }
+}
+
+impl ContextProvider for ContextProviderWithTasks {
+    fn associated_tasks(&self) -> Option<TaskTemplates> {
+        Some(self.templates.clone())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use gpui::TestAppContext;

crates/task/src/lib.rs

@@ -185,6 +185,20 @@ impl TaskVariables {
     pub fn extend(&mut self, other: Self) {
         self.0.extend(other.0);
     }
+    /// Get the value associated with given variable name, if there is one.
+    pub fn get(&self, key: &VariableName) -> Option<&str> {
+        self.0.get(key).map(|s| s.as_str())
+    }
+    /// Clear out variables obtained from tree-sitter queries, which are prefixed with '_' character
+    pub fn sweep(&mut self) {
+        self.0.retain(|name, _| {
+            if let VariableName::Custom(name) = name {
+                !name.starts_with('_')
+            } else {
+                true
+            }
+        })
+    }
 }
 
 impl FromIterator<(VariableName, String)> for TaskVariables {

crates/tasks_ui/src/lib.rs

@@ -153,8 +153,8 @@ mod tests {
 
     use editor::Editor;
     use gpui::{Entity, TestAppContext};
-    use language::{BasicContextProvider, Language, LanguageConfig};
+    use language::{Language, LanguageConfig};
-    use project::{FakeFs, Project};
+    use project::{BasicContextProvider, FakeFs, Project};
     use serde_json::json;
     use task::{TaskContext, TaskVariables, VariableName};
     use ui::VisualContext;
@@ -191,7 +191,7 @@ mod tests {
             }),
         )
         .await;
-
+        let project = Project::test(fs, ["/dir".as_ref()], cx).await;
         let rust_language = Arc::new(
             Language::new(
                 LanguageConfig::default(),
@@ -203,7 +203,7 @@ mod tests {
             name: (_) @name) @item"#,
             )
             .unwrap()
-            .with_context_provider(Some(Arc::new(BasicContextProvider))),
+            .with_context_provider(Some(Arc::new(BasicContextProvider::new(project.clone())))),
         );
 
         let typescript_language = Arc::new(
@@ -221,9 +221,9 @@ mod tests {
                       ")" @context)) @item"#,
             )
             .unwrap()
-            .with_context_provider(Some(Arc::new(BasicContextProvider))),
+            .with_context_provider(Some(Arc::new(BasicContextProvider::new(project.clone())))),
         );
-        let project = Project::test(fs, ["/dir".as_ref()], cx).await;
+
         let worktree_id = project.update(cx, |project, cx| {
             project.worktrees().next().unwrap().read(cx).id()
         });

crates/tasks_ui/src/modal.rs

@@ -540,8 +540,8 @@ mod tests {
 
     use editor::Editor;
     use gpui::{TestAppContext, VisualTestContext};
-    use language::{ContextProviderWithTasks, Language, LanguageConfig, LanguageMatcher, Point};
+    use language::{Language, LanguageConfig, LanguageMatcher, Point};
-    use project::{FakeFs, Project};
+    use project::{ContextProviderWithTasks, FakeFs, Project};
     use serde_json::json;
     use task::TaskTemplates;
     use workspace::CloseInactiveTabsAndPanes;