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 })
159 }
160
161 pub fn local(
162 buffer_store: WeakEntity<BufferStore>,
163 worktree_store: Entity<WorktreeStore>,
164 toolchain_store: Arc<dyn LanguageToolchainStore>,
165 environment: Entity<ProjectEnvironment>,
166 cx: &mut Context<Self>,
167 ) -> Self {
168 Self::Functional(StoreState {
169 mode: StoreMode::Local {
170 downstream_client: None,
171 environment,
172 },
173 task_inventory: Inventory::new(cx),
174 buffer_store,
175 toolchain_store,
176 worktree_store,
177 })
178 }
179
180 pub fn remote(
181 buffer_store: WeakEntity<BufferStore>,
182 worktree_store: Entity<WorktreeStore>,
183 toolchain_store: Arc<dyn LanguageToolchainStore>,
184 upstream_client: AnyProtoClient,
185 project_id: u64,
186 cx: &mut Context<Self>,
187 ) -> Self {
188 Self::Functional(StoreState {
189 mode: StoreMode::Remote {
190 upstream_client,
191 project_id,
192 },
193 task_inventory: Inventory::new(cx),
194 buffer_store,
195 toolchain_store,
196 worktree_store,
197 })
198 }
199
200 pub fn task_context_for_location(
201 &self,
202 captured_variables: TaskVariables,
203 location: Location,
204 cx: &mut App,
205 ) -> Task<Option<TaskContext>> {
206 match self {
207 TaskStore::Functional(state) => match &state.mode {
208 StoreMode::Local { environment, .. } => local_task_context_for_location(
209 state.worktree_store.clone(),
210 state.toolchain_store.clone(),
211 environment.clone(),
212 captured_variables,
213 location,
214 cx,
215 ),
216 StoreMode::Remote {
217 upstream_client,
218 project_id,
219 } => remote_task_context_for_location(
220 *project_id,
221 upstream_client.clone(),
222 state.worktree_store.clone(),
223 captured_variables,
224 location,
225 state.toolchain_store.clone(),
226 cx,
227 ),
228 },
229 TaskStore::Noop => Task::ready(None),
230 }
231 }
232
233 pub fn task_inventory(&self) -> Option<&Entity<Inventory>> {
234 match self {
235 TaskStore::Functional(state) => Some(&state.task_inventory),
236 TaskStore::Noop => None,
237 }
238 }
239
240 pub fn shared(&mut self, remote_id: u64, new_downstream_client: AnyProtoClient, _cx: &mut App) {
241 if let Self::Functional(StoreState {
242 mode: StoreMode::Local {
243 downstream_client, ..
244 },
245 ..
246 }) = self
247 {
248 *downstream_client = Some((new_downstream_client, remote_id));
249 }
250 }
251
252 pub fn unshared(&mut self, _: &mut Context<Self>) {
253 if let Self::Functional(StoreState {
254 mode: StoreMode::Local {
255 downstream_client, ..
256 },
257 ..
258 }) = self
259 {
260 *downstream_client = None;
261 }
262 }
263
264 pub(super) fn update_user_tasks(
265 &self,
266 location: TaskSettingsLocation<'_>,
267 raw_tasks_json: Option<&str>,
268 cx: &mut Context<Self>,
269 ) -> Result<(), InvalidSettingsError> {
270 let task_inventory = match self {
271 TaskStore::Functional(state) => &state.task_inventory,
272 TaskStore::Noop => return Ok(()),
273 };
274 let raw_tasks_json = raw_tasks_json
275 .map(|json| json.trim())
276 .filter(|json| !json.is_empty());
277
278 task_inventory.update(cx, |inventory, _| {
279 inventory.update_file_based_tasks(location, raw_tasks_json)
280 })
281 }
282
283 pub(super) fn update_user_debug_scenarios(
284 &self,
285 location: TaskSettingsLocation<'_>,
286 raw_tasks_json: Option<&str>,
287 cx: &mut Context<Self>,
288 ) -> Result<(), InvalidSettingsError> {
289 let task_inventory = match self {
290 TaskStore::Functional(state) => &state.task_inventory,
291 TaskStore::Noop => return Ok(()),
292 };
293 let raw_tasks_json = raw_tasks_json
294 .map(|json| json.trim())
295 .filter(|json| !json.is_empty());
296
297 task_inventory.update(cx, |inventory, _| {
298 inventory.update_file_based_scenarios(location, raw_tasks_json)
299 })
300 }
301}
302
303fn local_task_context_for_location(
304 worktree_store: Entity<WorktreeStore>,
305 toolchain_store: Arc<dyn LanguageToolchainStore>,
306 environment: Entity<ProjectEnvironment>,
307 captured_variables: TaskVariables,
308 location: Location,
309 cx: &App,
310) -> Task<Option<TaskContext>> {
311 let worktree_id = location.buffer.read(cx).file().map(|f| f.worktree_id(cx));
312 let worktree_abs_path = worktree_id
313 .and_then(|worktree_id| worktree_store.read(cx).worktree_for_id(worktree_id, cx))
314 .and_then(|worktree| worktree.read(cx).root_dir());
315 let fs = worktree_store.read(cx).fs();
316
317 cx.spawn(async move |cx| {
318 let project_env = environment
319 .update(cx, |environment, cx| {
320 environment.get_buffer_environment(&location.buffer, &worktree_store, cx)
321 })
322 .ok()?
323 .await;
324
325 let mut task_variables = cx
326 .update(|cx| {
327 combine_task_variables(
328 captured_variables,
329 fs,
330 worktree_store.clone(),
331 location,
332 project_env.clone(),
333 BasicContextProvider::new(worktree_store),
334 toolchain_store,
335 cx,
336 )
337 })
338 .ok()?
339 .await
340 .log_err()?;
341 // Remove all custom entries starting with _, as they're not intended for use by the end user.
342 task_variables.sweep();
343
344 Some(TaskContext {
345 project_env: project_env.unwrap_or_default(),
346 cwd: worktree_abs_path.map(|p| p.to_path_buf()),
347 task_variables,
348 })
349 })
350}
351
352fn remote_task_context_for_location(
353 project_id: u64,
354 upstream_client: AnyProtoClient,
355 worktree_store: Entity<WorktreeStore>,
356 captured_variables: TaskVariables,
357 location: Location,
358 toolchain_store: Arc<dyn LanguageToolchainStore>,
359 cx: &mut App,
360) -> Task<Option<TaskContext>> {
361 cx.spawn(async move |cx| {
362 // 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).
363 let mut remote_context = cx
364 .update(|cx| {
365 let worktree_root = worktree_root(&worktree_store, &location, cx);
366
367 BasicContextProvider::new(worktree_store).build_context(
368 &TaskVariables::default(),
369 ContextLocation {
370 fs: None,
371 worktree_root,
372 file_location: &location,
373 },
374 None,
375 toolchain_store,
376 cx,
377 )
378 })
379 .ok()?
380 .await
381 .log_err()
382 .unwrap_or_default();
383 remote_context.extend(captured_variables);
384
385 let buffer_id = cx
386 .update(|cx| location.buffer.read(cx).remote_id().to_proto())
387 .ok()?;
388 let context_task = upstream_client.request(proto::TaskContextForLocation {
389 project_id,
390 location: Some(proto::Location {
391 buffer_id,
392 start: Some(serialize_anchor(&location.range.start)),
393 end: Some(serialize_anchor(&location.range.end)),
394 }),
395 task_variables: remote_context
396 .into_iter()
397 .map(|(k, v)| (k.to_string(), v))
398 .collect(),
399 });
400 let task_context = context_task.await.log_err()?;
401 Some(TaskContext {
402 cwd: task_context.cwd.map(PathBuf::from),
403 task_variables: task_context
404 .task_variables
405 .into_iter()
406 .filter_map(
407 |(variable_name, variable_value)| match variable_name.parse() {
408 Ok(variable_name) => Some((variable_name, variable_value)),
409 Err(()) => {
410 log::error!("Unknown variable name: {variable_name}");
411 None
412 }
413 },
414 )
415 .collect(),
416 project_env: task_context.project_env.into_iter().collect(),
417 })
418 })
419}
420
421fn worktree_root(
422 worktree_store: &Entity<WorktreeStore>,
423 location: &Location,
424 cx: &mut App,
425) -> Option<PathBuf> {
426 location
427 .buffer
428 .read(cx)
429 .file()
430 .map(|f| f.worktree_id(cx))
431 .and_then(|worktree_id| worktree_store.read(cx).worktree_for_id(worktree_id, cx))
432 .and_then(|worktree| {
433 let worktree = worktree.read(cx);
434 if !worktree.is_visible() {
435 return None;
436 }
437 let root_entry = worktree.root_entry()?;
438 if !root_entry.is_dir() {
439 return None;
440 }
441 Some(worktree.absolutize(&root_entry.path))
442 })
443}
444
445fn combine_task_variables(
446 mut captured_variables: TaskVariables,
447 fs: Option<Arc<dyn Fs>>,
448 worktree_store: Entity<WorktreeStore>,
449 location: Location,
450 project_env: Option<HashMap<String, String>>,
451 baseline: BasicContextProvider,
452 toolchain_store: Arc<dyn LanguageToolchainStore>,
453 cx: &mut App,
454) -> Task<anyhow::Result<TaskVariables>> {
455 let language_context_provider = location
456 .buffer
457 .read(cx)
458 .language()
459 .and_then(|language| language.context_provider());
460 cx.spawn(async move |cx| {
461 let baseline = cx
462 .update(|cx| {
463 let worktree_root = worktree_root(&worktree_store, &location, cx);
464 baseline.build_context(
465 &captured_variables,
466 ContextLocation {
467 fs: fs.clone(),
468 worktree_root,
469 file_location: &location,
470 },
471 project_env.clone(),
472 toolchain_store.clone(),
473 cx,
474 )
475 })?
476 .await
477 .context("building basic default context")?;
478 captured_variables.extend(baseline);
479 if let Some(provider) = language_context_provider {
480 captured_variables.extend(
481 cx.update(|cx| {
482 let worktree_root = worktree_root(&worktree_store, &location, cx);
483 provider.build_context(
484 &captured_variables,
485 ContextLocation {
486 fs,
487 worktree_root,
488 file_location: &location,
489 },
490 project_env,
491 toolchain_store,
492 cx,
493 )
494 })?
495 .await
496 .context("building provider context")?,
497 );
498 }
499 Ok(captured_variables)
500 })
501}