task_store.rs

  1use std::{
  2    path::{Path, PathBuf},
  3    sync::Arc,
  4};
  5
  6use anyhow::Context as _;
  7use collections::HashMap;
  8use fs::Fs;
  9use gpui::{App, AsyncApp, Context, Entity, EventEmitter, Task, WeakEntity};
 10use language::{
 11    ContextLocation, ContextProvider as _, LanguageToolchainStore, Location,
 12    proto::{deserialize_anchor, serialize_anchor},
 13};
 14use rpc::{AnyProtoClient, TypedEnvelope, proto};
 15use settings::{InvalidSettingsError, SettingsLocation};
 16use task::{TaskContext, TaskVariables, VariableName};
 17use text::{BufferId, OffsetRangeExt};
 18use util::ResultExt;
 19
 20use crate::{
 21    BasicContextProvider, Inventory, ProjectEnvironment, buffer_store::BufferStore,
 22    worktree_store::WorktreeStore,
 23};
 24
 25// platform-dependent warning
 26pub enum TaskStore {
 27    Functional(StoreState),
 28    Noop,
 29}
 30
 31pub struct StoreState {
 32    mode: StoreMode,
 33    task_inventory: Entity<Inventory>,
 34    buffer_store: WeakEntity<BufferStore>,
 35    worktree_store: Entity<WorktreeStore>,
 36    toolchain_store: Arc<dyn LanguageToolchainStore>,
 37}
 38
 39enum StoreMode {
 40    Local {
 41        downstream_client: Option<(AnyProtoClient, u64)>,
 42        environment: Entity<ProjectEnvironment>,
 43    },
 44    Remote {
 45        upstream_client: AnyProtoClient,
 46        project_id: u64,
 47    },
 48}
 49
 50impl EventEmitter<crate::Event> for TaskStore {}
 51
 52#[derive(Debug)]
 53pub enum TaskSettingsLocation<'a> {
 54    Global(&'a Path),
 55    Worktree(SettingsLocation<'a>),
 56}
 57
 58impl TaskStore {
 59    pub fn init(client: Option<&AnyProtoClient>) {
 60        if let Some(client) = client {
 61            client.add_entity_request_handler(Self::handle_task_context_for_location);
 62        }
 63    }
 64
 65    async fn handle_task_context_for_location(
 66        store: Entity<Self>,
 67        envelope: TypedEnvelope<proto::TaskContextForLocation>,
 68        mut cx: AsyncApp,
 69    ) -> anyhow::Result<proto::TaskContext> {
 70        let location = envelope
 71            .payload
 72            .location
 73            .context("no location given for task context handling")?;
 74        let (buffer_store, is_remote) = store.read_with(&cx, |store, _| {
 75            Ok(match store {
 76                TaskStore::Functional(state) => (
 77                    state.buffer_store.clone(),
 78                    match &state.mode {
 79                        StoreMode::Local { .. } => false,
 80                        StoreMode::Remote { .. } => true,
 81                    },
 82                ),
 83                TaskStore::Noop => {
 84                    anyhow::bail!("empty task store cannot handle task context requests")
 85                }
 86            })
 87        })??;
 88        let buffer_store = buffer_store
 89            .upgrade()
 90            .context("no buffer store when handling task context request")?;
 91
 92        let buffer_id = BufferId::new(location.buffer_id).with_context(|| {
 93            format!(
 94                "cannot handle task context request for invalid buffer id: {}",
 95                location.buffer_id
 96            )
 97        })?;
 98
 99        let start = location
100            .start
101            .and_then(deserialize_anchor)
102            .context("missing task context location start")?;
103        let end = location
104            .end
105            .and_then(deserialize_anchor)
106            .context("missing task context location end")?;
107        let buffer = buffer_store
108            .update(&mut cx, |buffer_store, cx| {
109                if is_remote {
110                    buffer_store.wait_for_remote_buffer(buffer_id, cx)
111                } else {
112                    Task::ready(
113                        buffer_store
114                            .get(buffer_id)
115                            .with_context(|| format!("no local buffer with id {buffer_id}")),
116                    )
117                }
118            })?
119            .await?;
120
121        let location = Location {
122            buffer,
123            range: start..end,
124        };
125        let context_task = store.update(&mut cx, |store, cx| {
126            let captured_variables = {
127                let mut variables = TaskVariables::from_iter(
128                    envelope
129                        .payload
130                        .task_variables
131                        .into_iter()
132                        .filter_map(|(k, v)| Some((k.parse().log_err()?, v))),
133                );
134
135                let snapshot = location.buffer.read(cx).snapshot();
136                let range = location.range.to_offset(&snapshot);
137
138                for range in snapshot.runnable_ranges(range) {
139                    for (capture_name, value) in range.extra_captures {
140                        variables.insert(VariableName::Custom(capture_name.into()), value);
141                    }
142                }
143                variables
144            };
145            store.task_context_for_location(captured_variables, location, cx)
146        })?;
147        let task_context = context_task.await.unwrap_or_default();
148        Ok(proto::TaskContext {
149            project_env: task_context.project_env.into_iter().collect(),
150            cwd: task_context
151                .cwd
152                .map(|cwd| cwd.to_string_lossy().to_string()),
153            task_variables: task_context
154                .task_variables
155                .into_iter()
156                .map(|(variable_name, variable_value)| (variable_name.to_string(), variable_value))
157                .collect(),
158        })
159    }
160
161    pub fn local(
162        fs: Arc<dyn Fs>,
163        buffer_store: WeakEntity<BufferStore>,
164        worktree_store: Entity<WorktreeStore>,
165        toolchain_store: Arc<dyn LanguageToolchainStore>,
166        environment: Entity<ProjectEnvironment>,
167        cx: &mut Context<Self>,
168    ) -> Self {
169        Self::Functional(StoreState {
170            mode: StoreMode::Local {
171                downstream_client: None,
172                environment,
173            },
174            task_inventory: Inventory::new(fs, cx),
175            buffer_store,
176            toolchain_store,
177            worktree_store,
178        })
179    }
180
181    pub fn remote(
182        fs: Arc<dyn Fs>,
183        buffer_store: WeakEntity<BufferStore>,
184        worktree_store: Entity<WorktreeStore>,
185        toolchain_store: Arc<dyn LanguageToolchainStore>,
186        upstream_client: AnyProtoClient,
187        project_id: u64,
188        cx: &mut Context<Self>,
189    ) -> Self {
190        Self::Functional(StoreState {
191            mode: StoreMode::Remote {
192                upstream_client,
193                project_id,
194            },
195            task_inventory: Inventory::new(fs, cx),
196            buffer_store,
197            toolchain_store,
198            worktree_store,
199        })
200    }
201
202    pub fn task_context_for_location(
203        &self,
204        captured_variables: TaskVariables,
205        location: Location,
206        cx: &mut App,
207    ) -> Task<Option<TaskContext>> {
208        match self {
209            TaskStore::Functional(state) => match &state.mode {
210                StoreMode::Local { environment, .. } => local_task_context_for_location(
211                    state.worktree_store.clone(),
212                    state.toolchain_store.clone(),
213                    environment.clone(),
214                    captured_variables,
215                    location,
216                    cx,
217                ),
218                StoreMode::Remote {
219                    upstream_client,
220                    project_id,
221                } => remote_task_context_for_location(
222                    *project_id,
223                    upstream_client.clone(),
224                    state.worktree_store.clone(),
225                    captured_variables,
226                    location,
227                    state.toolchain_store.clone(),
228                    cx,
229                ),
230            },
231            TaskStore::Noop => Task::ready(None),
232        }
233    }
234
235    pub fn task_inventory(&self) -> Option<&Entity<Inventory>> {
236        match self {
237            TaskStore::Functional(state) => Some(&state.task_inventory),
238            TaskStore::Noop => None,
239        }
240    }
241
242    pub fn shared(&mut self, remote_id: u64, new_downstream_client: AnyProtoClient, _cx: &mut App) {
243        if let Self::Functional(StoreState {
244            mode: StoreMode::Local {
245                downstream_client, ..
246            },
247            ..
248        }) = self
249        {
250            *downstream_client = Some((new_downstream_client, remote_id));
251        }
252    }
253
254    pub fn unshared(&mut self, _: &mut Context<Self>) {
255        if let Self::Functional(StoreState {
256            mode: StoreMode::Local {
257                downstream_client, ..
258            },
259            ..
260        }) = self
261        {
262            *downstream_client = None;
263        }
264    }
265
266    pub(super) fn update_user_tasks(
267        &self,
268        location: TaskSettingsLocation<'_>,
269        raw_tasks_json: Option<&str>,
270        cx: &mut Context<Self>,
271    ) -> Result<(), InvalidSettingsError> {
272        let task_inventory = match self {
273            TaskStore::Functional(state) => &state.task_inventory,
274            TaskStore::Noop => return Ok(()),
275        };
276        let raw_tasks_json = raw_tasks_json
277            .map(|json| json.trim())
278            .filter(|json| !json.is_empty());
279
280        task_inventory.update(cx, |inventory, _| {
281            inventory.update_file_based_tasks(location, raw_tasks_json)
282        })
283    }
284
285    pub(super) fn update_user_debug_scenarios(
286        &self,
287        location: TaskSettingsLocation<'_>,
288        raw_tasks_json: Option<&str>,
289        cx: &mut Context<Self>,
290    ) -> Result<(), InvalidSettingsError> {
291        let task_inventory = match self {
292            TaskStore::Functional(state) => &state.task_inventory,
293            TaskStore::Noop => return Ok(()),
294        };
295        let raw_tasks_json = raw_tasks_json
296            .map(|json| json.trim())
297            .filter(|json| !json.is_empty());
298
299        task_inventory.update(cx, |inventory, _| {
300            inventory.update_file_based_scenarios(location, raw_tasks_json)
301        })
302    }
303}
304
305fn local_task_context_for_location(
306    worktree_store: Entity<WorktreeStore>,
307    toolchain_store: Arc<dyn LanguageToolchainStore>,
308    environment: Entity<ProjectEnvironment>,
309    captured_variables: TaskVariables,
310    location: Location,
311    cx: &App,
312) -> Task<Option<TaskContext>> {
313    let worktree_id = location.buffer.read(cx).file().map(|f| f.worktree_id(cx));
314    let worktree_abs_path = worktree_id
315        .and_then(|worktree_id| worktree_store.read(cx).worktree_for_id(worktree_id, cx))
316        .and_then(|worktree| worktree.read(cx).root_dir());
317    let fs = worktree_store.read(cx).fs();
318
319    cx.spawn(async move |cx| {
320        let project_env = environment
321            .update(cx, |environment, cx| {
322                environment.get_buffer_environment(&location.buffer, &worktree_store, cx)
323            })
324            .ok()?
325            .await;
326
327        let mut task_variables = cx
328            .update(|cx| {
329                combine_task_variables(
330                    captured_variables,
331                    fs,
332                    worktree_store.clone(),
333                    location,
334                    project_env.clone(),
335                    BasicContextProvider::new(worktree_store),
336                    toolchain_store,
337                    cx,
338                )
339            })
340            .ok()?
341            .await
342            .log_err()?;
343        // Remove all custom entries starting with _, as they're not intended for use by the end user.
344        task_variables.sweep();
345
346        Some(TaskContext {
347            project_env: project_env.unwrap_or_default(),
348            cwd: worktree_abs_path.map(|p| p.to_path_buf()),
349            task_variables,
350        })
351    })
352}
353
354fn remote_task_context_for_location(
355    project_id: u64,
356    upstream_client: AnyProtoClient,
357    worktree_store: Entity<WorktreeStore>,
358    captured_variables: TaskVariables,
359    location: Location,
360    toolchain_store: Arc<dyn LanguageToolchainStore>,
361    cx: &mut App,
362) -> Task<Option<TaskContext>> {
363    cx.spawn(async move |cx| {
364        // We need to gather a client context, as the headless one may lack certain information (e.g. tree-sitter parsing is disabled there, so symbols are not available).
365        let mut remote_context = cx
366            .update(|cx| {
367                let worktree_root = worktree_root(&worktree_store, &location, cx);
368
369                BasicContextProvider::new(worktree_store).build_context(
370                    &TaskVariables::default(),
371                    ContextLocation {
372                        fs: None,
373                        worktree_root,
374                        file_location: &location,
375                    },
376                    None,
377                    toolchain_store,
378                    cx,
379                )
380            })
381            .ok()?
382            .await
383            .log_err()
384            .unwrap_or_default();
385        remote_context.extend(captured_variables);
386
387        let buffer_id = cx
388            .update(|cx| location.buffer.read(cx).remote_id().to_proto())
389            .ok()?;
390        let context_task = upstream_client.request(proto::TaskContextForLocation {
391            project_id,
392            location: Some(proto::Location {
393                buffer_id,
394                start: Some(serialize_anchor(&location.range.start)),
395                end: Some(serialize_anchor(&location.range.end)),
396            }),
397            task_variables: remote_context
398                .into_iter()
399                .map(|(k, v)| (k.to_string(), v))
400                .collect(),
401        });
402        let task_context = context_task.await.log_err()?;
403        Some(TaskContext {
404            cwd: task_context.cwd.map(PathBuf::from),
405            task_variables: task_context
406                .task_variables
407                .into_iter()
408                .filter_map(
409                    |(variable_name, variable_value)| match variable_name.parse() {
410                        Ok(variable_name) => Some((variable_name, variable_value)),
411                        Err(()) => {
412                            log::error!("Unknown variable name: {variable_name}");
413                            None
414                        }
415                    },
416                )
417                .collect(),
418            project_env: task_context.project_env.into_iter().collect(),
419        })
420    })
421}
422
423fn worktree_root(
424    worktree_store: &Entity<WorktreeStore>,
425    location: &Location,
426    cx: &mut App,
427) -> Option<PathBuf> {
428    location
429        .buffer
430        .read(cx)
431        .file()
432        .map(|f| f.worktree_id(cx))
433        .and_then(|worktree_id| worktree_store.read(cx).worktree_for_id(worktree_id, cx))
434        .and_then(|worktree| {
435            let worktree = worktree.read(cx);
436            if !worktree.is_visible() {
437                return None;
438            }
439            let root_entry = worktree.root_entry()?;
440            if !root_entry.is_dir() {
441                return None;
442            }
443            worktree.absolutize(&root_entry.path).ok()
444        })
445}
446
447fn combine_task_variables(
448    mut captured_variables: TaskVariables,
449    fs: Option<Arc<dyn Fs>>,
450    worktree_store: Entity<WorktreeStore>,
451    location: Location,
452    project_env: Option<HashMap<String, String>>,
453    baseline: BasicContextProvider,
454    toolchain_store: Arc<dyn LanguageToolchainStore>,
455    cx: &mut App,
456) -> Task<anyhow::Result<TaskVariables>> {
457    let language_context_provider = location
458        .buffer
459        .read(cx)
460        .language()
461        .and_then(|language| language.context_provider());
462    cx.spawn(async move |cx| {
463        let baseline = cx
464            .update(|cx| {
465                let worktree_root = worktree_root(&worktree_store, &location, cx);
466                baseline.build_context(
467                    &captured_variables,
468                    ContextLocation {
469                        fs: fs.clone(),
470                        worktree_root,
471                        file_location: &location,
472                    },
473                    project_env.clone(),
474                    toolchain_store.clone(),
475                    cx,
476                )
477            })?
478            .await
479            .context("building basic default context")?;
480        captured_variables.extend(baseline);
481        if let Some(provider) = language_context_provider {
482            captured_variables.extend(
483                cx.update(|cx| {
484                    let worktree_root = worktree_root(&worktree_store, &location, cx);
485                    provider.build_context(
486                        &captured_variables,
487                        ContextLocation {
488                            fs,
489                            worktree_root,
490                            file_location: &location,
491                        },
492                        project_env,
493                        toolchain_store,
494                        cx,
495                    )
496                })?
497                .await
498                .context("building provider context")?,
499            );
500        }
501        Ok(captured_variables)
502    })
503}