1use crate::{
2 example::{Example, ExampleContext},
3 headless::EpAppState,
4 load_project::run_load_project,
5 progress::{InfoStyle, Progress, Step, StepProgress},
6};
7use anyhow::Context as _;
8use collections::HashSet;
9use edit_prediction::{DebugEvent, EditPredictionStore};
10use futures::{FutureExt as _, StreamExt as _, channel::mpsc};
11use gpui::{AsyncApp, Entity};
12use language::Buffer;
13use project::Project;
14use std::sync::Arc;
15use std::time::Duration;
16
17pub async fn run_context_retrieval(
18 example: &mut Example,
19 app_state: Arc<EpAppState>,
20 mut cx: AsyncApp,
21) -> anyhow::Result<()> {
22 if example.context.is_some() {
23 return Ok(());
24 }
25
26 run_load_project(example, app_state.clone(), cx.clone()).await?;
27
28 let step_progress: Arc<StepProgress> = Progress::global()
29 .start(Step::Context, &example.name)
30 .into();
31
32 let state = example.state.as_ref().unwrap();
33 let project = state.project.clone();
34
35 let _lsp_handle = project.update(&mut cx, |project, cx| {
36 project.register_buffer_with_language_servers(&state.buffer, cx)
37 })?;
38 wait_for_language_servers_to_start(&project, &state.buffer, &step_progress, &mut cx).await?;
39
40 let ep_store = cx.update(|cx| {
41 EditPredictionStore::try_global(cx).context("EditPredictionStore not initialized")
42 })??;
43
44 let mut events = ep_store.update(&mut cx, |store, cx| {
45 store.register_buffer(&state.buffer, &project, cx);
46 store.set_use_context(true);
47 store.refresh_context(&project, &state.buffer, state.cursor_position, cx);
48 store.debug_info(&project, cx)
49 })?;
50
51 while let Some(event) = events.next().await {
52 match event {
53 DebugEvent::ContextRetrievalFinished(_) => {
54 break;
55 }
56 _ => {}
57 }
58 }
59
60 let context_files =
61 ep_store.update(&mut cx, |store, cx| store.context_for_project(&project, cx))?;
62
63 let excerpt_count: usize = context_files.iter().map(|f| f.excerpts.len()).sum();
64 step_progress.set_info(format!("{} excerpts", excerpt_count), InfoStyle::Normal);
65
66 example.context = Some(ExampleContext {
67 files: context_files,
68 });
69 Ok(())
70}
71
72async fn wait_for_language_servers_to_start(
73 project: &Entity<Project>,
74 buffer: &Entity<Buffer>,
75 step_progress: &Arc<StepProgress>,
76 cx: &mut AsyncApp,
77) -> anyhow::Result<()> {
78 let lsp_store = project.read_with(cx, |project, _| project.lsp_store())?;
79
80 let (language_server_ids, mut starting_language_server_ids) = buffer
81 .update(cx, |buffer, cx| {
82 lsp_store.update(cx, |lsp_store, cx| {
83 let ids = lsp_store.language_servers_for_local_buffer(buffer, cx);
84 let starting_ids = ids
85 .iter()
86 .copied()
87 .filter(|id| !lsp_store.language_server_statuses.contains_key(&id))
88 .collect::<HashSet<_>>();
89 (ids, starting_ids)
90 })
91 })
92 .unwrap_or_default();
93
94 step_progress.set_substatus(format!("waiting for {} LSPs", language_server_ids.len()));
95
96 let timeout = cx
97 .background_executor()
98 .timer(Duration::from_secs(60 * 5))
99 .shared();
100
101 let (mut tx, mut rx) = mpsc::channel(language_server_ids.len());
102 let added_subscription = cx.subscribe(project, {
103 let step_progress = step_progress.clone();
104 move |_, event, _| match event {
105 project::Event::LanguageServerAdded(language_server_id, name, _) => {
106 step_progress.set_substatus(format!("LSP started: {}", name));
107 tx.try_send(*language_server_id).ok();
108 }
109 _ => {}
110 }
111 });
112
113 while !starting_language_server_ids.is_empty() {
114 futures::select! {
115 language_server_id = rx.next() => {
116 if let Some(id) = language_server_id {
117 starting_language_server_ids.remove(&id);
118 }
119 },
120 _ = timeout.clone().fuse() => {
121 return Err(anyhow::anyhow!("LSP wait timed out after 5 minutes"));
122 }
123 }
124 }
125
126 drop(added_subscription);
127
128 if !language_server_ids.is_empty() {
129 project
130 .update(cx, |project, cx| project.save_buffer(buffer.clone(), cx))?
131 .detach();
132 }
133
134 let (mut tx, mut rx) = mpsc::channel(language_server_ids.len());
135 let subscriptions = [
136 cx.subscribe(&lsp_store, {
137 let step_progress = step_progress.clone();
138 move |_, event, _| {
139 if let project::LspStoreEvent::LanguageServerUpdate {
140 message:
141 client::proto::update_language_server::Variant::WorkProgress(
142 client::proto::LspWorkProgress {
143 message: Some(message),
144 ..
145 },
146 ),
147 ..
148 } = event
149 {
150 step_progress.set_substatus(message.clone());
151 }
152 }
153 }),
154 cx.subscribe(project, {
155 let step_progress = step_progress.clone();
156 move |_, event, cx| match event {
157 project::Event::DiskBasedDiagnosticsFinished { language_server_id } => {
158 let lsp_store = lsp_store.read(cx);
159 let name = lsp_store
160 .language_server_adapter_for_id(*language_server_id)
161 .unwrap()
162 .name();
163 step_progress.set_substatus(format!("LSP idle: {}", name));
164 tx.try_send(*language_server_id).ok();
165 }
166 _ => {}
167 }
168 }),
169 ];
170
171 project
172 .update(cx, |project, cx| project.save_buffer(buffer.clone(), cx))?
173 .await?;
174
175 let mut pending_language_server_ids = HashSet::from_iter(language_server_ids.into_iter());
176 while !pending_language_server_ids.is_empty() {
177 futures::select! {
178 language_server_id = rx.next() => {
179 if let Some(id) = language_server_id {
180 pending_language_server_ids.remove(&id);
181 }
182 },
183 _ = timeout.clone().fuse() => {
184 return Err(anyhow::anyhow!("LSP wait timed out after 5 minutes"));
185 }
186 }
187 }
188
189 drop(subscriptions);
190 step_progress.clear_substatus();
191 Ok(())
192}