tenant.openagents/omega

1use std::sync::Arc;
2use std::time::Duration;
3
4use crate::Editor;
5use collections::{HashMap, HashSet};
6use futures::TryFutureExt;
7use gpui::{App, Entity, Task};
8use gpui::{AsyncApp, FutureExt};
9use language::Buffer;
10use language::Language;
11use lsp::LanguageServerId;
12use lsp::LanguageServerName;
13use project::LanguageServerToQuery;
14use project::LocationLink;
15use project::Project;
16use project::TaskSourceKind;
17use project::lsp_store::lsp_ext_command::GetLspRunnables;
18use task::ResolvedTask;
19use task::TaskContext;
20use text::BufferId;
21use ui::SharedString;
22use util::ResultExt as _;
23
24pub(crate) fn find_specific_language_server_in_selection<F>(
25    editor: &Editor,
26    cx: &mut App,
27    filter_language: F,
28    language_server_name: LanguageServerName,
29) -> Option<(
30    text::Anchor,
31    Arc<Language>,
32    LanguageServerId,
33    Entity<Buffer>,
34)>
35where
36    F: Fn(&Language) -> bool,
37{
38    let project = editor.project.clone()?;
39    let multi_buffer = editor.buffer();
40    let mut seen_buffer_ids = HashSet::default();
41    editor
42        .selections
43        .disjoint_anchors_arc()
44        .iter()
45        .find_map(|selection| {
46            let multi_buffer = multi_buffer.read(cx);
47            let multi_buffer_snapshot = multi_buffer.snapshot(cx);
48            let (position, buffer) = multi_buffer_snapshot
49                .anchor_to_buffer_anchor(selection.head())
50                .and_then(|(anchor, _)| Some((anchor, multi_buffer.buffer(anchor.buffer_id)?)))?;
51            if !seen_buffer_ids.insert(buffer.read(cx).remote_id()) {
52                return None;
53            }
54
55            let language = buffer.read(cx).language_at(position)?;
56            if filter_language(&language) {
57                let server_id = buffer.update(cx, |buffer, cx| {
58                    project
59                        .read(cx)
60                        .language_server_id_for_name(buffer, &language_server_name, cx)
61                })?;
62                Some((position, language, server_id, buffer))
63            } else {
64                None
65            }
66        })
67}
68
69async fn lsp_task_context(
70    project: &Entity<Project>,
71    buffer: &Entity<Buffer>,
72    cx: &mut AsyncApp,
73) -> Option<TaskContext> {
74    let (worktree_store, environment) = project.read_with(cx, |project, _| {
75        (project.worktree_store(), project.environment().clone())
76    });
77
78    let worktree_abs_path = cx.update(|cx| {
79        let worktree_id = buffer.read(cx).file().map(|f| f.worktree_id(cx));
80
81        worktree_id
82            .and_then(|worktree_id| worktree_store.read(cx).worktree_for_id(worktree_id, cx))
83            .and_then(|worktree| worktree.read(cx).root_dir())
84    });
85
86    let project_env = environment
87        .update(cx, |environment, cx| {
88            environment.buffer_environment(buffer, &worktree_store, cx)
89        })
90        .await;
91
92    Some(TaskContext {
93        cwd: worktree_abs_path.map(|p| p.to_path_buf()),
94        project_env: project_env.unwrap_or_default(),
95        ..TaskContext::default()
96    })
97}
98
99pub fn lsp_tasks(
100    project: Entity<Project>,
101    task_sources: &HashMap<LanguageServerName, Vec<BufferId>>,
102    for_position: Option<text::Anchor>,
103    cx: &mut App,
104) -> Task<Vec<(TaskSourceKind, Vec<(Option<LocationLink>, ResolvedTask)>)>> {
105    let lsp_task_sources = task_sources
106        .iter()
107        .filter_map(|(name, buffer_ids)| {
108            let buffers = buffer_ids
109                .iter()
110                .filter(|&&buffer_id| match for_position {
111                    Some(for_position) => for_position.buffer_id == buffer_id,
112                    None => true,
113                })
114                .filter_map(|&buffer_id| project.read(cx).buffer_for_id(buffer_id, cx))
115                .collect::<Vec<_>>();
116
117            let server_id = buffers.iter().find_map(|buffer| {
118                project.read_with(cx, |project, cx| {
119                    project.language_server_id_for_name(buffer.read(cx), name, cx)
120                })
121            });
122            server_id.zip(Some(buffers))
123        })
124        .collect::<Vec<_>>();
125
126    cx.spawn(async move |cx| {
127        cx.spawn(async move |cx| {
128            let mut lsp_tasks = HashMap::default();
129            for (server_id, buffers) in lsp_task_sources {
130                let mut new_lsp_tasks = Vec::new();
131                for buffer in buffers {
132                    let source_kind = match buffer.update(cx, |buffer, _| {
133                        buffer.language().map(|language| language.name())
134                    }) {
135                        Some(language_name) => TaskSourceKind::Lsp {
136                            server: server_id,
137                            language_name: SharedString::from(language_name),
138                        },
139                        None => continue,
140                    };
141                    let id_base = source_kind.to_id_base();
142                    let lsp_buffer_context = lsp_task_context(&project, &buffer, cx)
143                        .await
144                        .unwrap_or_default();
145
146                    let runnables_task = project.update(cx, |project, cx| {
147                        let buffer_id = buffer.read(cx).remote_id();
148                        project.request_lsp(
149                            buffer,
150                            LanguageServerToQuery::Other(server_id),
151                            GetLspRunnables {
152                                buffer_id,
153                                position: for_position,
154                            },
155                            cx,
156                        )
157                    });
158                    if let Some(new_runnables) = runnables_task.await.log_err() {
159                        new_lsp_tasks.extend(new_runnables.runnables.into_iter().filter_map(
160                            |(location, runnable)| {
161                                let resolved_task =
162                                    runnable.resolve_task(&id_base, &lsp_buffer_context)?;
163                                Some((location, resolved_task))
164                            },
165                        ));
166                    }
167                    if !new_lsp_tasks.is_empty() {
168                        lsp_tasks
169                            .entry(source_kind)
170                            .or_insert_with(Vec::new)
171                            .append(&mut new_lsp_tasks);
172                    }
173                }
174            }
175            lsp_tasks.into_iter().collect()
176        })
177        .with_timeout(Duration::from_millis(200), &cx.background_executor())
178        .unwrap_or_else(|_| {
179            log::debug!("Timed out waiting for LSP tasks");
180            Vec::new()
181        })
182        .await
183    })
184}
185
186#[cfg(test)]
187mod tests {
188    use std::sync::Arc;
189
190    use futures::StreamExt as _;
191    use gpui::{AppContext as _, Entity, TestAppContext};
192    use language::{FakeLspAdapter, Language};
193    use languages::rust_lang;
194    use lsp::{LanguageServerId, LanguageServerName};
195    use multi_buffer::MultiBuffer;
196    use project::{FakeFs, Project};
197    use util::path;
198
199    use crate::{MoveToEnd, editor_tests::init_test, test::build_editor_with_project};
200
201    use super::find_specific_language_server_in_selection;
202
203    #[gpui::test]
204    async fn test_find_language_server_at_end_of_file(cx: &mut TestAppContext) {
205        init_test(cx, |_| {});
206
207        let fs = FakeFs::new(cx.executor());
208        fs.insert_file(path!("/file.rs"), "fn main() {}".into())
209            .await;
210
211        let project = Project::test(fs, [path!("/file.rs").as_ref()], cx).await;
212        let language_registry = project.read_with(cx, |project, _| project.languages().clone());
213        language_registry.add(rust_lang());
214        let mut fake_servers =
215            language_registry.register_fake_lsp("Rust", FakeLspAdapter::default());
216
217        let underlying_buffer = project
218            .update(cx, |project, cx| {
219                project.open_local_buffer(path!("/file.rs"), cx)
220            })
221            .await
222            .unwrap();
223
224        let buffer = cx.new(|cx| MultiBuffer::singleton(underlying_buffer.clone(), cx));
225        let (editor, cx) = cx.add_window_view(|window, cx| {
226            build_editor_with_project(project.clone(), buffer, window, cx)
227        });
228
229        let fake_server = fake_servers.next().await.unwrap();
230        cx.executor().run_until_parked();
231
232        let expected_server_id = fake_server.server.server_id();
233        let language_server_name = LanguageServerName::new_static("the-fake-language-server");
234        let filter = |language: &Language| language.name().as_ref() == "Rust";
235
236        let assert_result = |result: Option<(
237            text::Anchor,
238            Arc<Language>,
239            LanguageServerId,
240            Entity<language::Buffer>,
241        )>,
242                             message: &str| {
243            let (_, language, server_id, buffer) = result.expect(message);
244            assert_eq!(
245                language.name().as_ref(),
246                "Rust",
247                "{message}: wrong language"
248            );
249            assert_eq!(server_id, expected_server_id, "{message}: wrong server ID");
250            assert_eq!(buffer, underlying_buffer, "{message}: wrong buffer");
251        };
252
253        editor.update(cx, |editor, cx| {
254            assert_result(
255                find_specific_language_server_in_selection(
256                    editor,
257                    cx,
258                    filter,
259                    language_server_name.clone(),
260                ),
261                "should find correct language server at beginning of file",
262            );
263        });
264
265        editor.update_in(cx, |editor, window, cx| {
266            editor.move_to_end(&MoveToEnd, window, cx);
267        });
268
269        editor.update(cx, |editor, cx| {
270            assert_result(
271                find_specific_language_server_in_selection(
272                    editor,
273                    cx,
274                    filter,
275                    language_server_name.clone(),
276                ),
277                "should find correct language server at end of file",
278            );
279        });
280    }
281}
282