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().into_owned()),
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 is_windows: cfg!(windows),
159 })
160 }
161
162 pub fn local(
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(cx),
175 buffer_store,
176 toolchain_store,
177 worktree_store,
178 })
179 }
180
181 pub fn remote(
182 buffer_store: WeakEntity<BufferStore>,
183 worktree_store: Entity<WorktreeStore>,
184 toolchain_store: Arc<dyn LanguageToolchainStore>,
185 upstream_client: AnyProtoClient,
186 project_id: u64,
187 cx: &mut Context<Self>,
188 ) -> Self {
189 Self::Functional(StoreState {
190 mode: StoreMode::Remote {
191 upstream_client,
192 project_id,
193 },
194 task_inventory: Inventory::new(cx),
195 buffer_store,
196 toolchain_store,
197 worktree_store,
198 })
199 }
200
201 pub fn task_context_for_location(
202 &self,
203 captured_variables: TaskVariables,
204 location: Location,
205 cx: &mut App,
206 ) -> Task<Option<TaskContext>> {
207 match self {
208 TaskStore::Functional(state) => match &state.mode {
209 StoreMode::Local { environment, .. } => local_task_context_for_location(
210 state.worktree_store.clone(),
211 state.toolchain_store.clone(),
212 environment.clone(),
213 captured_variables,
214 location,
215 cx,
216 ),
217 StoreMode::Remote {
218 upstream_client,
219 project_id,
220 } => remote_task_context_for_location(
221 *project_id,
222 upstream_client.clone(),
223 state.worktree_store.clone(),
224 captured_variables,
225 location,
226 state.toolchain_store.clone(),
227 cx,
228 ),
229 },
230 TaskStore::Noop => Task::ready(None),
231 }
232 }
233
234 pub fn task_inventory(&self) -> Option<&Entity<Inventory>> {
235 match self {
236 TaskStore::Functional(state) => Some(&state.task_inventory),
237 TaskStore::Noop => None,
238 }
239 }
240
241 pub fn shared(&mut self, remote_id: u64, new_downstream_client: AnyProtoClient, _cx: &mut App) {
242 if let Self::Functional(StoreState {
243 mode: StoreMode::Local {
244 downstream_client, ..
245 },
246 ..
247 }) = self
248 {
249 *downstream_client = Some((new_downstream_client, remote_id));
250 }
251 }
252
253 pub fn unshared(&mut self, _: &mut Context<Self>) {
254 if let Self::Functional(StoreState {
255 mode: StoreMode::Local {
256 downstream_client, ..
257 },
258 ..
259 }) = self
260 {
261 *downstream_client = None;
262 }
263 }
264
265 pub(super) fn update_user_tasks(
266 &self,
267 location: TaskSettingsLocation<'_>,
268 raw_tasks_json: Option<&str>,
269 cx: &mut Context<Self>,
270 ) -> Result<(), InvalidSettingsError> {
271 let task_inventory = match self {
272 TaskStore::Functional(state) => &state.task_inventory,
273 TaskStore::Noop => return Ok(()),
274 };
275 let raw_tasks_json = raw_tasks_json
276 .map(|json| json.trim())
277 .filter(|json| !json.is_empty());
278
279 task_inventory.update(cx, |inventory, _| {
280 inventory.update_file_based_tasks(location, raw_tasks_json)
281 })
282 }
283
284 pub(super) fn update_user_debug_scenarios(
285 &self,
286 location: TaskSettingsLocation<'_>,
287 raw_tasks_json: Option<&str>,
288 cx: &mut Context<Self>,
289 ) -> Result<(), InvalidSettingsError> {
290 let task_inventory = match self {
291 TaskStore::Functional(state) => &state.task_inventory,
292 TaskStore::Noop => return Ok(()),
293 };
294 let raw_tasks_json = raw_tasks_json
295 .map(|json| json.trim())
296 .filter(|json| !json.is_empty());
297
298 task_inventory.update(cx, |inventory, _| {
299 inventory.update_file_based_scenarios(location, raw_tasks_json)
300 })
301 }
302}
303
304fn local_task_context_for_location(
305 worktree_store: Entity<WorktreeStore>,
306 toolchain_store: Arc<dyn LanguageToolchainStore>,
307 environment: Entity<ProjectEnvironment>,
308 captured_variables: TaskVariables,
309 location: Location,
310 cx: &App,
311) -> Task<Option<TaskContext>> {
312 let worktree_id = location.buffer.read(cx).file().map(|f| f.worktree_id(cx));
313 let worktree_abs_path = worktree_id
314 .and_then(|worktree_id| worktree_store.read(cx).worktree_for_id(worktree_id, cx))
315 .and_then(|worktree| worktree.read(cx).root_dir());
316 let fs = worktree_store.read(cx).fs();
317
318 cx.spawn(async move |cx| {
319 let project_env = environment
320 .update(cx, |environment, cx| {
321 environment.get_buffer_environment(&location.buffer, &worktree_store, cx)
322 })
323 .ok()?
324 .await;
325
326 let mut task_variables = cx
327 .update(|cx| {
328 combine_task_variables(
329 captured_variables,
330 fs,
331 worktree_store.clone(),
332 location,
333 project_env.clone(),
334 BasicContextProvider::new(worktree_store),
335 toolchain_store,
336 cx,
337 )
338 })
339 .ok()?
340 .await
341 .log_err()?;
342 // Remove all custom entries starting with _, as they're not intended for use by the end user.
343 task_variables.sweep();
344
345 Some(TaskContext {
346 project_env: project_env.unwrap_or_default(),
347 cwd: worktree_abs_path.map(|p| p.to_path_buf()),
348 task_variables,
349 is_windows: cfg!(windows),
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 is_windows: task_context.is_windows,
420 })
421 })
422}
423
424fn worktree_root(
425 worktree_store: &Entity<WorktreeStore>,
426 location: &Location,
427 cx: &mut App,
428) -> Option<PathBuf> {
429 location
430 .buffer
431 .read(cx)
432 .file()
433 .map(|f| f.worktree_id(cx))
434 .and_then(|worktree_id| worktree_store.read(cx).worktree_for_id(worktree_id, cx))
435 .and_then(|worktree| {
436 let worktree = worktree.read(cx);
437 if !worktree.is_visible() {
438 return None;
439 }
440 let root_entry = worktree.root_entry()?;
441 if !root_entry.is_dir() {
442 return None;
443 }
444 Some(worktree.absolutize(&root_entry.path))
445 })
446}
447
448fn combine_task_variables(
449 mut captured_variables: TaskVariables,
450 fs: Option<Arc<dyn Fs>>,
451 worktree_store: Entity<WorktreeStore>,
452 location: Location,
453 project_env: Option<HashMap<String, String>>,
454 baseline: BasicContextProvider,
455 toolchain_store: Arc<dyn LanguageToolchainStore>,
456 cx: &mut App,
457) -> Task<anyhow::Result<TaskVariables>> {
458 let language_context_provider = location
459 .buffer
460 .read(cx)
461 .language()
462 .and_then(|language| language.context_provider());
463 cx.spawn(async move |cx| {
464 let baseline = cx
465 .update(|cx| {
466 let worktree_root = worktree_root(&worktree_store, &location, cx);
467 baseline.build_context(
468 &captured_variables,
469 ContextLocation {
470 fs: fs.clone(),
471 worktree_root,
472 file_location: &location,
473 },
474 project_env.clone(),
475 toolchain_store.clone(),
476 cx,
477 )
478 })?
479 .await
480 .context("building basic default context")?;
481 captured_variables.extend(baseline);
482 if let Some(provider) = language_context_provider {
483 captured_variables.extend(
484 cx.update(|cx| {
485 let worktree_root = worktree_root(&worktree_store, &location, cx);
486 provider.build_context(
487 &captured_variables,
488 ContextLocation {
489 fs,
490 worktree_root,
491 file_location: &location,
492 },
493 project_env,
494 toolchain_store,
495 cx,
496 )
497 })?
498 .await
499 .context("building provider context")?,
500 );
501 }
502 Ok(captured_variables)
503 })
504}