tenant.openagents/omega

1use collections::HashMap;
2use futures::channel::mpsc::UnboundedReceiver;
3use language::{Language, LanguageRegistry};
4use lsp::{
5    FakeLanguageServer, LanguageServerBinary, TextDocumentSyncCapability, TextDocumentSyncKind, Uri,
6};
7use parking_lot::Mutex;
8use project::Fs;
9use std::{ops::Range, path::PathBuf, sync::Arc};
10use tree_sitter::{Parser, QueryCursor, StreamingIterator, Tree};
11
12/// Registers a fake language server that implements go-to-definition and
13/// go-to-type-definition using tree-sitter, making the assumption that all
14/// names are unique, and all variables' types are explicitly declared.
15pub fn register_fake_definition_server(
16    language_registry: &Arc<LanguageRegistry>,
17    language: Arc<Language>,
18    fs: Arc<dyn Fs>,
19) -> UnboundedReceiver<FakeLanguageServer> {
20    let index = Arc::new(Mutex::new(DefinitionIndex::new(language.clone())));
21
22    language_registry.register_fake_lsp(
23        language.name(),
24        language::FakeLspAdapter {
25            name: "fake-definition-lsp",
26            initialization_options: None,
27            prettier_plugins: Vec::new(),
28            disk_based_diagnostics_progress_token: None,
29            disk_based_diagnostics_sources: Vec::new(),
30            language_server_binary: LanguageServerBinary {
31                path: PathBuf::from("fake-definition-lsp"),
32                arguments: Vec::new(),
33                env: None,
34            },
35            capabilities: lsp::ServerCapabilities {
36                definition_provider: Some(lsp::OneOf::Left(true)),
37                type_definition_provider: Some(lsp::TypeDefinitionProviderCapability::Simple(true)),
38                text_document_sync: Some(TextDocumentSyncCapability::Kind(
39                    TextDocumentSyncKind::FULL,
40                )),
41                ..Default::default()
42            },
43            label_for_completion: None,
44            initializer: Some(Box::new({
45                move |server| {
46                    server.handle_notification::<lsp::notification::DidOpenTextDocument, _>({
47                        let index = index.clone();
48                        move |params, _cx| {
49                            index
50                                .lock()
51                                .open_buffer(params.text_document.uri, &params.text_document.text);
52                        }
53                    });
54
55                    server.handle_notification::<lsp::notification::DidCloseTextDocument, _>({
56                        let index = index.clone();
57                        let fs = fs.clone();
58                        move |params, cx| {
59                            let uri = params.text_document.uri;
60                            let path = uri.to_file_path().ok();
61                            index.lock().mark_buffer_closed(&uri);
62
63                            if let Some(path) = path {
64                                let index = index.clone();
65                                let fs = fs.clone();
66                                cx.spawn(async move |_cx| {
67                                    if let Ok(content) = fs.load(&path).await {
68                                        index.lock().index_file(uri, &content);
69                                    }
70                                })
71                                .detach();
72                            }
73                        }
74                    });
75
76                    server.handle_notification::<lsp::notification::DidChangeWatchedFiles, _>({
77                        let index = index.clone();
78                        let fs = fs.clone();
79                        move |params, cx| {
80                            let index = index.clone();
81                            let fs = fs.clone();
82                            cx.spawn(async move |_cx| {
83                                for event in params.changes {
84                                    if index.lock().is_buffer_open(&event.uri) {
85                                        continue;
86                                    }
87
88                                    match event.typ {
89                                        lsp::FileChangeType::DELETED => {
90                                            index.lock().remove_definitions_for_file(&event.uri);
91                                        }
92                                        lsp::FileChangeType::CREATED
93                                        | lsp::FileChangeType::CHANGED => {
94                                            if let Some(path) = event.uri.to_file_path().ok() {
95                                                if let Ok(content) = fs.load(&path).await {
96                                                    index.lock().index_file(event.uri, &content);
97                                                }
98                                            }
99                                        }
100                                        _ => {}
101                                    }
102                                }
103                            })
104                            .detach();
105                        }
106                    });
107
108                    server.handle_notification::<lsp::notification::DidChangeTextDocument, _>({
109                        let index = index.clone();
110                        move |params, _cx| {
111                            if let Some(change) = params.content_changes.into_iter().last() {
112                                index
113                                    .lock()
114                                    .index_file(params.text_document.uri, &change.text);
115                            }
116                        }
117                    });
118
119                    server.handle_notification::<lsp::notification::DidChangeWorkspaceFolders, _>(
120                        {
121                            let index = index.clone();
122                            let fs = fs.clone();
123                            move |params, cx| {
124                                let index = index.clone();
125                                let fs = fs.clone();
126                                let files = fs.as_fake().files();
127                                cx.spawn(async move |_cx| {
128                                    for folder in params.event.added {
129                                        let Ok(path) = folder.uri.to_file_path() else {
130                                            continue;
131                                        };
132                                        for file in &files {
133                                            if let Some(uri) = Uri::from_file_path(&file).ok()
134                                                && file.starts_with(&path)
135                                                && let Ok(content) = fs.load(&file).await
136                                            {
137                                                index.lock().index_file(uri, &content);
138                                            }
139                                        }
140                                    }
141                                })
142                                .detach();
143                            }
144                        },
145                    );
146
147                    server.set_request_handler::<lsp::request::GotoDefinition, _, _>({
148                        let index = index.clone();
149                        move |params, _cx| {
150                            let result = index.lock().get_definitions(
151                                params.text_document_position_params.text_document.uri,
152                                params.text_document_position_params.position,
153                            );
154                            async move { Ok(result) }
155                        }
156                    });
157
158                    server.set_request_handler::<lsp::request::GotoTypeDefinition, _, _>({
159                        let index = index.clone();
160                        move |params, _cx| {
161                            let result = index.lock().get_type_definitions(
162                                params.text_document_position_params.text_document.uri,
163                                params.text_document_position_params.position,
164                            );
165                            async move { Ok(result) }
166                        }
167                    });
168                }
169            })),
170        },
171    )
172}
173
174struct DefinitionIndex {
175    language: Arc<Language>,
176    definitions: HashMap<String, Vec<lsp::Location>>,
177    type_annotations_by_file: HashMap<Uri, HashMap<String, String>>,
178    files: HashMap<Uri, FileEntry>,
179}
180
181#[derive(Debug)]
182struct FileEntry {
183    contents: String,
184    is_open_in_buffer: bool,
185}
186
187impl DefinitionIndex {
188    fn new(language: Arc<Language>) -> Self {
189        Self {
190            language,
191            definitions: HashMap::default(),
192            type_annotations_by_file: HashMap::default(),
193            files: HashMap::default(),
194        }
195    }
196
197    fn remove_definitions_for_file(&mut self, uri: &Uri) {
198        self.definitions.retain(|_, locations| {
199            locations.retain(|loc| &loc.uri != uri);
200            !locations.is_empty()
201        });
202        self.type_annotations_by_file.remove(uri);
203        self.files.remove(uri);
204    }
205
206    fn open_buffer(&mut self, uri: Uri, content: &str) {
207        self.index_file_inner(uri, content, true);
208    }
209
210    fn mark_buffer_closed(&mut self, uri: &Uri) {
211        if let Some(entry) = self.files.get_mut(uri) {
212            entry.is_open_in_buffer = false;
213        }
214    }
215
216    fn is_buffer_open(&self, uri: &Uri) -> bool {
217        self.files
218            .get(uri)
219            .map(|entry| entry.is_open_in_buffer)
220            .unwrap_or(false)
221    }
222
223    fn index_file(&mut self, uri: Uri, content: &str) {
224        self.index_file_inner(uri, content, false);
225    }
226
227    fn index_file_inner(&mut self, uri: Uri, content: &str, is_open_in_buffer: bool) -> Option<()> {
228        self.remove_definitions_for_file(&uri);
229        let grammar = self.language.grammar()?;
230        let outline_config = grammar.outline_config.as_ref()?;
231        let mut parser = Parser::new();
232        parser.set_language(&grammar.ts_language).ok()?;
233        let tree = parser.parse(content, None)?;
234        let declarations = extract_declarations_from_tree(&tree, content, outline_config);
235        for (name, byte_range) in declarations {
236            let range = byte_range_to_lsp_range(content, byte_range);
237            let location = lsp::Location {
238                uri: uri.clone(),
239                range,
240            };
241            self.definitions
242                .entry(name)
243                .or_insert_with(Vec::new)
244                .push(location);
245        }
246        for (name, location) in extract_extra_definition_locations(content) {
247            self.definitions
248                .entry(name)
249                .or_insert_with(Vec::new)
250                .push(location);
251        }
252
253        let type_annotations = extract_type_annotations(content)
254            .into_iter()
255            .collect::<HashMap<_, _>>();
256        self.type_annotations_by_file
257            .insert(uri.clone(), type_annotations);
258
259        self.files.insert(
260            uri,
261            FileEntry {
262                contents: content.to_string(),
263                is_open_in_buffer,
264            },
265        );
266
267        Some(())
268    }
269
270    fn get_definitions(
271        &mut self,
272        uri: Uri,
273        position: lsp::Position,
274    ) -> Option<lsp::GotoDefinitionResponse> {
275        let entry = self.files.get(&uri)?;
276        let name = word_at_position(&entry.contents, position)?;
277        let locations = self.definitions.get(name).cloned()?;
278        Some(lsp::GotoDefinitionResponse::Array(locations))
279    }
280
281    fn get_type_definitions(
282        &mut self,
283        uri: Uri,
284        position: lsp::Position,
285    ) -> Option<lsp::GotoDefinitionResponse> {
286        let entry = self.files.get(&uri)?;
287        let name = word_at_position(&entry.contents, position)?;
288
289        if let Some(type_name) = self
290            .type_annotations_by_file
291            .get(&uri)
292            .and_then(|annotations| annotations.get(name))
293        {
294            if let Some(locations) = self.definitions.get(type_name) {
295                return Some(lsp::GotoDefinitionResponse::Array(locations.clone()));
296            }
297        }
298
299        // If the identifier itself is an uppercase name (a type), return its own definition.
300        // This mirrors real LSP behavior where GotoTypeDefinition on a type name
301        // resolves to that type's definition.
302        if name.starts_with(|c: char| c.is_uppercase()) {
303            if let Some(locations) = self.definitions.get(name) {
304                return Some(lsp::GotoDefinitionResponse::Array(locations.clone()));
305            }
306        }
307
308        None
309    }
310}
311
312fn extract_extra_definition_locations(content: &str) -> Vec<(String, lsp::Location)> {
313    content
314        .lines()
315        .filter_map(|line| {
316            let mut parts = line
317                .trim()
318                .strip_prefix("// fake-definition-lsp-extra ")?
319                .split_whitespace();
320            let name = parts.next()?.to_string();
321            let path = PathBuf::from(parts.next()?);
322            let start_row = parts.next()?.parse().ok()?;
323            let start_column = parts.next()?.parse().ok()?;
324            let end_row = parts.next()?.parse().ok()?;
325            let end_column = parts.next()?.parse().ok()?;
326            Some((
327                name,
328                lsp::Location::new(
329                    Uri::from_file_path(path).ok()?,
330                    lsp::Range::new(
331                        lsp::Position::new(start_row, start_column),
332                        lsp::Position::new(end_row, end_column),
333                    ),
334                ),
335            ))
336        })
337        .collect()
338}
339
340/// Extracts `identifier_name -> type_name` mappings from field declarations
341/// and function parameters. For example, `owner: Arc<Person>` produces
342/// `"owner" -> "Person"` by unwrapping common generic wrappers.
343fn extract_type_annotations(content: &str) -> Vec<(String, String)> {
344    let mut annotations = Vec::new();
345    for line in content.lines() {
346        let trimmed = line.trim();
347        if trimmed.starts_with("//")
348            || trimmed.starts_with("use ")
349            || trimmed.starts_with("pub use ")
350        {
351            continue;
352        }
353
354        let Some(colon_idx) = trimmed.find(':') else {
355            continue;
356        };
357
358        // The part before `:` should end with an identifier name.
359        let left = trimmed[..colon_idx].trim();
360        let Some(name) = left.split_whitespace().last() else {
361            continue;
362        };
363
364        if name.is_empty() || !name.chars().all(|c| c.is_alphanumeric() || c == '_') {
365            continue;
366        }
367
368        // Skip names that start uppercase — they're type names, not variables/fields.
369        if name.starts_with(|c: char| c.is_uppercase()) {
370            continue;
371        }
372
373        let right = trimmed[colon_idx + 1..].trim();
374        let type_name = extract_base_type_name(right);
375
376        if !type_name.is_empty() && type_name.starts_with(|c: char| c.is_uppercase()) {
377            annotations.push((name.to_string(), type_name));
378        }
379    }
380    annotations
381}
382
383/// Unwraps common generic wrappers (Arc, Box, Rc, Option, Vec) and trait
384/// object prefixes (dyn, impl) to find the concrete type name. For example:
385/// `Arc<Person>` → `"Person"`, `Box<dyn Trait>` → `"Trait"`.
386fn extract_base_type_name(type_str: &str) -> String {
387    let trimmed = type_str
388        .trim()
389        .trim_start_matches('&')
390        .trim_start_matches("mut ")
391        .trim_end_matches(',')
392        .trim_end_matches('{')
393        .trim_end_matches(')')
394        .trim()
395        .trim_start_matches("dyn ")
396        .trim_start_matches("impl ")
397        .trim();
398
399    if let Some(angle_start) = trimmed.find('<') {
400        let outer = &trimmed[..angle_start];
401        if matches!(outer, "Arc" | "Box" | "Rc" | "Option" | "Vec" | "Cow") {
402            let inner_end = trimmed.rfind('>').unwrap_or(trimmed.len());
403            let inner = &trimmed[angle_start + 1..inner_end];
404            return extract_base_type_name(inner);
405        }
406        return outer.to_string();
407    }
408
409    if let Some(call_start) = trimmed.find("::") {
410        let outer = &trimmed[..call_start];
411        if matches!(outer, "Arc" | "Box" | "Rc" | "Option" | "Vec" | "Cow") {
412            let rest = trimmed[call_start + 2..].trim_start();
413            if let Some(paren_start) = rest.find('(') {
414                let inner = &rest[paren_start + 1..];
415                let inner = inner.trim();
416                if !inner.is_empty() {
417                    return extract_base_type_name(inner);
418                }
419            }
420        }
421    }
422
423    trimmed
424        .split(|c: char| !c.is_alphanumeric() && c != '_')
425        .next()
426        .unwrap_or("")
427        .to_string()
428}
429
430fn extract_declarations_from_tree(
431    tree: &Tree,
432    content: &str,
433    outline_config: &language::OutlineConfig,
434) -> Vec<(String, Range<usize>)> {
435    let mut cursor = QueryCursor::new();
436    let mut declarations = Vec::new();
437    let mut matches = cursor.matches(&outline_config.query, tree.root_node(), content.as_bytes());
438    while let Some(query_match) = matches.next() {
439        let mut name_range: Option<Range<usize>> = None;
440        let mut has_item_range = false;
441
442        for capture in query_match.captures {
443            let range = capture.node.byte_range();
444            if capture.index == outline_config.name_capture_ix {
445                name_range = Some(range);
446            } else if capture.index == outline_config.item_capture_ix {
447                has_item_range = true;
448            }
449        }
450
451        if let Some(name_range) = name_range
452            && has_item_range
453        {
454            let name = content[name_range.clone()].to_string();
455            if declarations.iter().any(|(n, _)| n == &name) {
456                continue;
457            }
458            declarations.push((name, name_range));
459        }
460    }
461    declarations
462}
463
464fn byte_range_to_lsp_range(content: &str, byte_range: Range<usize>) -> lsp::Range {
465    let start = byte_offset_to_position(content, byte_range.start);
466    let end = byte_offset_to_position(content, byte_range.end);
467    lsp::Range { start, end }
468}
469
470fn byte_offset_to_position(content: &str, offset: usize) -> lsp::Position {
471    let mut line = 0;
472    let mut character = 0;
473    let mut current_offset = 0;
474    for ch in content.chars() {
475        if current_offset >= offset {
476            break;
477        }
478        if ch == '\n' {
479            line += 1;
480            character = 0;
481        } else {
482            character += 1;
483        }
484        current_offset += ch.len_utf8();
485    }
486    lsp::Position { line, character }
487}
488
489fn word_at_position(content: &str, position: lsp::Position) -> Option<&str> {
490    let mut lines = content.lines();
491    let line = lines.nth(position.line as usize)?;
492    let column = position.character as usize;
493    if column > line.len() {
494        return None;
495    }
496    let start = line[..column]
497        .rfind(|c: char| !c.is_alphanumeric() && c != '_')
498        .map(|i| i + 1)
499        .unwrap_or(0);
500    let end = line[column..]
501        .find(|c: char| !c.is_alphanumeric() && c != '_')
502        .map(|i| i + column)
503        .unwrap_or(line.len());
504    Some(&line[start..end]).filter(|word| !word.is_empty())
505}
506