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