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.buffer_environment(&location.buffer, &worktree_store, cx)
321 })
322 .await;
323
324 let mut task_variables = cx
325 .update(|cx| {
326 combine_task_variables(
327 captured_variables,
328 fs,
329 worktree_store.clone(),
330 location,
331 project_env.clone(),
332 BasicContextProvider::new(worktree_store),
333 toolchain_store,
334 cx,
335 )
336 })
337 .await
338 .log_err()?;
339 // Remove all custom entries starting with _, as they're not intended for use by the end user.
340 task_variables.sweep();
341
342 Some(TaskContext {
343 project_env: project_env.unwrap_or_default(),
344 cwd: worktree_abs_path.map(|p| p.to_path_buf()),
345 task_variables,
346 })
347 })
348}
349
350fn remote_task_context_for_location(
351 project_id: u64,
352 upstream_client: AnyProtoClient,
353 worktree_store: Entity<WorktreeStore>,
354 captured_variables: TaskVariables,
355 location: Location,
356 toolchain_store: Arc<dyn LanguageToolchainStore>,
357 cx: &mut App,
358) -> Task<Option<TaskContext>> {
359 cx.spawn(async move |cx| {
360 // 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).
361 let mut remote_context = cx
362 .update(|cx| {
363 let worktree_root = worktree_root(&worktree_store, &location, cx);
364
365 BasicContextProvider::new(worktree_store).build_context(
366 &TaskVariables::default(),
367 ContextLocation {
368 fs: None,
369 worktree_root,
370 file_location: &location,
371 },
372 None,
373 toolchain_store,
374 cx,
375 )
376 })
377 .await
378 .log_err()
379 .unwrap_or_default();
380 remote_context.extend(captured_variables);
381
382 let buffer_id = cx.update(|cx| location.buffer.read(cx).remote_id().to_proto());
383 let context_task = upstream_client.request(proto::TaskContextForLocation {
384 project_id,
385 location: Some(proto::Location {
386 buffer_id,
387 start: Some(serialize_anchor(&location.range.start)),
388 end: Some(serialize_anchor(&location.range.end)),
389 }),
390 task_variables: remote_context
391 .into_iter()
392 .map(|(k, v)| (k.to_string(), v))
393 .collect(),
394 });
395 let task_context = context_task.await.log_err()?;
396 Some(TaskContext {
397 cwd: task_context.cwd.map(PathBuf::from),
398 task_variables: task_context
399 .task_variables
400 .into_iter()
401 .filter_map(
402 |(variable_name, variable_value)| match variable_name.parse() {
403 Ok(variable_name) => Some((variable_name, variable_value)),
404 Err(()) => {
405 log::error!("Unknown variable name: {variable_name}");
406 None
407 }
408 },
409 )
410 .collect(),
411 project_env: task_context.project_env.into_iter().collect(),
412 })
413 })
414}
415
416fn worktree_root(
417 worktree_store: &Entity<WorktreeStore>,
418 location: &Location,
419 cx: &mut App,
420) -> Option<PathBuf> {
421 location
422 .buffer
423 .read(cx)
424 .file()
425 .map(|f| f.worktree_id(cx))
426 .and_then(|worktree_id| worktree_store.read(cx).worktree_for_id(worktree_id, cx))
427 .and_then(|worktree| {
428 let worktree = worktree.read(cx);
429 if !worktree.is_visible() {
430 return None;
431 }
432 let root_entry = worktree.root_entry()?;
433 if !root_entry.is_dir() {
434 return None;
435 }
436 Some(worktree.absolutize(&root_entry.path))
437 })
438}
439
440fn combine_task_variables(
441 mut captured_variables: TaskVariables,
442 fs: Option<Arc<dyn Fs>>,
443 worktree_store: Entity<WorktreeStore>,
444 location: Location,
445 project_env: Option<HashMap<String, String>>,
446 baseline: BasicContextProvider,
447 toolchain_store: Arc<dyn LanguageToolchainStore>,
448 cx: &mut App,
449) -> Task<anyhow::Result<TaskVariables>> {
450 let language_context_provider = location
451 .buffer
452 .read(cx)
453 .language()
454 .and_then(|language| language.context_provider());
455 cx.spawn(async move |cx| {
456 let baseline = cx
457 .update(|cx| {
458 let worktree_root = worktree_root(&worktree_store, &location, cx);
459 baseline.build_context(
460 &captured_variables,
461 ContextLocation {
462 fs: fs.clone(),
463 worktree_root,
464 file_location: &location,
465 },
466 project_env.clone(),
467 toolchain_store.clone(),
468 cx,
469 )
470 })
471 .await
472 .context("building basic default context")?;
473 captured_variables.extend(baseline);
474 if let Some(provider) = language_context_provider {
475 captured_variables.extend(
476 cx.update(|cx| {
477 let worktree_root = worktree_root(&worktree_store, &location, cx);
478 provider.build_context(
479 &captured_variables,
480 ContextLocation {
481 fs,
482 worktree_root,
483 file_location: &location,
484 },
485 project_env,
486 toolchain_store,
487 cx,
488 )
489 })
490 .await
491 .context("building provider context")?,
492 );
493 }
494 Ok(captured_variables)
495 })
496}