Rework inlay hints system

crates/agent_ui/src/acp/message_editor.rs

@@ -11,10 +11,10 @@ use assistant_slash_commands::codeblock_fence_for_path;
 use collections::{HashMap, HashSet};
 use editor::{
     Addon, Anchor, AnchorRangeExt, ContextMenuOptions, ContextMenuPlacement, Editor, EditorElement,
-    EditorEvent, EditorMode, EditorSnapshot, EditorStyle, ExcerptId, FoldPlaceholder, InlayId,
+    EditorEvent, EditorMode, EditorSnapshot, EditorStyle, ExcerptId, FoldPlaceholder, Inlay,
     MultiBuffer, ToOffset,
     actions::Paste,
-    display_map::{Crease, CreaseId, FoldId, Inlay},
+    display_map::{Crease, CreaseId, FoldId},
 };
 use futures::{
     FutureExt as _,
@@ -29,7 +29,8 @@ use language::{Buffer, Language, language_settings::InlayHintKind};
 use language_model::LanguageModelImage;
 use postage::stream::Stream as _;
 use project::{
-    CompletionIntent, InlayHint, InlayHintLabel, Project, ProjectItem, ProjectPath, Worktree,
+    CompletionIntent, InlayHint, InlayHintLabel, InlayId, Project, ProjectItem, ProjectPath,
+    Worktree,
 };
 use prompt_store::{PromptId, PromptStore};
 use rope::Point;
@@ -75,7 +76,7 @@ pub enum MessageEditorEvent {
 
 impl EventEmitter<MessageEditorEvent> for MessageEditor {}
 
-const COMMAND_HINT_INLAY_ID: u32 = 0;
+const COMMAND_HINT_INLAY_ID: InlayId = InlayId::Hint(0);
 
 impl MessageEditor {
     pub fn new(
@@ -151,7 +152,7 @@ impl MessageEditor {
                         let has_new_hint = !new_hints.is_empty();
                         editor.splice_inlays(
                             if has_hint {
-                                &[InlayId::Hint(COMMAND_HINT_INLAY_ID)]
+                                &[COMMAND_HINT_INLAY_ID]
                             } else {
                                 &[]
                             },

crates/collab/src/rpc.rs

@@ -343,7 +343,6 @@ impl Server {
             .add_request_handler(forward_read_only_project_request::<proto::OpenBufferForSymbol>)
             .add_request_handler(forward_read_only_project_request::<proto::OpenBufferById>)
             .add_request_handler(forward_read_only_project_request::<proto::SynchronizeBuffers>)
-            .add_request_handler(forward_read_only_project_request::<proto::InlayHints>)
             .add_request_handler(forward_read_only_project_request::<proto::ResolveInlayHint>)
             .add_request_handler(forward_read_only_project_request::<proto::GetColorPresentation>)
             .add_request_handler(forward_read_only_project_request::<proto::OpenBufferByPath>)

crates/collab/src/tests/editor_tests.rs

@@ -1849,10 +1849,40 @@ async fn test_mutual_editor_inlay_hint_cache_update(
         ..lsp::ServerCapabilities::default()
     };
     client_a.language_registry().add(rust_lang());
+
+    // Set up the language server to return an additional inlay hint on each request.
+    let edits_made = Arc::new(AtomicUsize::new(0));
+    let closure_edits_made = Arc::clone(&edits_made);
     let mut fake_language_servers = client_a.language_registry().register_fake_lsp(
         "Rust",
         FakeLspAdapter {
             capabilities: capabilities.clone(),
+            initializer: Some(Box::new(move |fake_language_server| {
+                let closure_edits_made = closure_edits_made.clone();
+                fake_language_server.set_request_handler::<lsp::request::InlayHintRequest, _, _>(
+                    move |params, _| {
+                        let edits_made_2 = Arc::clone(&closure_edits_made);
+                        async move {
+                            assert_eq!(
+                                params.text_document.uri,
+                                lsp::Uri::from_file_path(path!("/a/main.rs")).unwrap(),
+                            );
+                            let edits_made =
+                                AtomicUsize::load(&edits_made_2, atomic::Ordering::Acquire);
+                            Ok(Some(vec![lsp::InlayHint {
+                                position: lsp::Position::new(0, edits_made as u32),
+                                label: lsp::InlayHintLabel::String(edits_made.to_string()),
+                                kind: None,
+                                text_edits: None,
+                                tooltip: None,
+                                padding_left: None,
+                                padding_right: None,
+                                data: None,
+                            }]))
+                        }
+                    },
+                );
+            })),
             ..FakeLspAdapter::default()
         },
     );
@@ -1894,61 +1924,20 @@ async fn test_mutual_editor_inlay_hint_cache_update(
         .unwrap();
 
     let (workspace_a, cx_a) = client_a.build_workspace(&project_a, cx_a);
-    executor.start_waiting();
 
     // The host opens a rust file.
-    let _buffer_a = project_a
-        .update(cx_a, |project, cx| {
-            project.open_local_buffer(path!("/a/main.rs"), cx)
-        })
-        .await
-        .unwrap();
-    let editor_a = workspace_a
-        .update_in(cx_a, |workspace, window, cx| {
-            workspace.open_path((worktree_id, rel_path("main.rs")), None, true, window, cx)
-        })
-        .await
-        .unwrap()
-        .downcast::<Editor>()
-        .unwrap();
-
+    let file_a = workspace_a.update_in(cx_a, |workspace, window, cx| {
+        workspace.open_path((worktree_id, rel_path("main.rs")), None, true, window, cx)
+    });
     let fake_language_server = fake_language_servers.next().await.unwrap();
-
-    // Set up the language server to return an additional inlay hint on each request.
-    let edits_made = Arc::new(AtomicUsize::new(0));
-    let closure_edits_made = Arc::clone(&edits_made);
-    fake_language_server
-        .set_request_handler::<lsp::request::InlayHintRequest, _, _>(move |params, _| {
-            let edits_made_2 = Arc::clone(&closure_edits_made);
-            async move {
-                assert_eq!(
-                    params.text_document.uri,
-                    lsp::Uri::from_file_path(path!("/a/main.rs")).unwrap(),
-                );
-                let edits_made = AtomicUsize::load(&edits_made_2, atomic::Ordering::Acquire);
-                Ok(Some(vec![lsp::InlayHint {
-                    position: lsp::Position::new(0, edits_made as u32),
-                    label: lsp::InlayHintLabel::String(edits_made.to_string()),
-                    kind: None,
-                    text_edits: None,
-                    tooltip: None,
-                    padding_left: None,
-                    padding_right: None,
-                    data: None,
-                }]))
-            }
-        })
-        .next()
-        .await
-        .unwrap();
-
+    let editor_a = file_a.await.unwrap().downcast::<Editor>().unwrap();
     executor.run_until_parked();
 
     let initial_edit = edits_made.load(atomic::Ordering::Acquire);
-    editor_a.update(cx_a, |editor, _| {
+    editor_a.update(cx_a, |editor, cx| {
         assert_eq!(
             vec![initial_edit.to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
             "Host should get its first hints when opens an editor"
         );
     });
@@ -1963,10 +1952,10 @@ async fn test_mutual_editor_inlay_hint_cache_update(
         .unwrap();
 
     executor.run_until_parked();
-    editor_b.update(cx_b, |editor, _| {
+    editor_b.update(cx_b, |editor, cx| {
         assert_eq!(
             vec![initial_edit.to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
             "Client should get its first hints when opens an editor"
         );
     });
@@ -1981,16 +1970,16 @@ async fn test_mutual_editor_inlay_hint_cache_update(
     cx_b.focus(&editor_b);
 
     executor.run_until_parked();
-    editor_a.update(cx_a, |editor, _| {
+    editor_a.update(cx_a, |editor, cx| {
         assert_eq!(
             vec![after_client_edit.to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
         );
     });
-    editor_b.update(cx_b, |editor, _| {
+    editor_b.update(cx_b, |editor, cx| {
         assert_eq!(
             vec![after_client_edit.to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
         );
     });
 
@@ -2004,16 +1993,16 @@ async fn test_mutual_editor_inlay_hint_cache_update(
     cx_a.focus(&editor_a);
 
     executor.run_until_parked();
-    editor_a.update(cx_a, |editor, _| {
+    editor_a.update(cx_a, |editor, cx| {
         assert_eq!(
             vec![after_host_edit.to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
         );
     });
-    editor_b.update(cx_b, |editor, _| {
+    editor_b.update(cx_b, |editor, cx| {
         assert_eq!(
             vec![after_host_edit.to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
         );
     });
 
@@ -2025,26 +2014,22 @@ async fn test_mutual_editor_inlay_hint_cache_update(
         .expect("inlay refresh request failed");
 
     executor.run_until_parked();
-    editor_a.update(cx_a, |editor, _| {
+    editor_a.update(cx_a, |editor, cx| {
         assert_eq!(
             vec![after_special_edit_for_refresh.to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
             "Host should react to /refresh LSP request"
         );
     });
-    editor_b.update(cx_b, |editor, _| {
+    editor_b.update(cx_b, |editor, cx| {
         assert_eq!(
             vec![after_special_edit_for_refresh.to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
             "Guest should get a /refresh LSP request propagated by host"
         );
     });
 }
 
-// This test started hanging on seed 2 after the theme settings
-// PR. The hypothesis is that it's been buggy for a while, but got lucky
-// on seeds.
-#[ignore]
 #[gpui::test(iterations = 10)]
 async fn test_inlay_hint_refresh_is_forwarded(
     cx_a: &mut TestAppContext,
@@ -2206,18 +2191,18 @@ async fn test_inlay_hint_refresh_is_forwarded(
     executor.finish_waiting();
 
     executor.run_until_parked();
-    editor_a.update(cx_a, |editor, _| {
+    editor_a.update(cx_a, |editor, cx| {
         assert!(
-            extract_hint_labels(editor).is_empty(),
+            extract_hint_labels(editor, cx).is_empty(),
             "Host should get no hints due to them turned off"
         );
     });
 
     executor.run_until_parked();
-    editor_b.update(cx_b, |editor, _| {
+    editor_b.update(cx_b, |editor, cx| {
         assert_eq!(
             vec!["initial hint".to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
             "Client should get its first hints when opens an editor"
         );
     });
@@ -2229,18 +2214,18 @@ async fn test_inlay_hint_refresh_is_forwarded(
         .into_response()
         .expect("inlay refresh request failed");
     executor.run_until_parked();
-    editor_a.update(cx_a, |editor, _| {
+    editor_a.update(cx_a, |editor, cx| {
         assert!(
-            extract_hint_labels(editor).is_empty(),
+            extract_hint_labels(editor, cx).is_empty(),
             "Host should get no hints due to them turned off, even after the /refresh"
         );
     });
 
     executor.run_until_parked();
-    editor_b.update(cx_b, |editor, _| {
+    editor_b.update(cx_b, |editor, cx| {
         assert_eq!(
             vec!["other hint".to_string()],
-            extract_hint_labels(editor),
+            extract_hint_labels(editor, cx),
             "Guest should get a /refresh LSP request propagated by host despite host hints are off"
         );
     });
@@ -4217,15 +4202,35 @@ fn tab_undo_assert(
     cx_b.assert_editor_state(expected_initial);
 }
 
-fn extract_hint_labels(editor: &Editor) -> Vec<String> {
-    let mut labels = Vec::new();
-    for hint in editor.inlay_hint_cache().hints() {
-        match hint.label {
-            project::InlayHintLabel::String(s) => labels.push(s),
-            _ => unreachable!(),
-        }
+fn extract_hint_labels(editor: &Editor, cx: &mut App) -> Vec<String> {
+    let lsp_store = editor.project().unwrap().read(cx).lsp_store();
+
+    let mut all_cached_labels = Vec::new();
+    let mut all_fetched_hints = Vec::new();
+    for buffer in editor.buffer().read(cx).all_buffers() {
+        lsp_store.update(cx, |lsp_store, cx| {
+            let hints = &lsp_store.latest_lsp_data(&buffer, cx).inlay_hints();
+            all_cached_labels.extend(hints.all_cached_hints().into_iter().map(|hint| {
+                let mut label = hint.text().to_string();
+                if hint.padding_left {
+                    label.insert(0, ' ');
+                }
+                if hint.padding_right {
+                    label.push_str(" ");
+                }
+                label
+            }));
+            all_fetched_hints.extend(hints.all_fetched_hints());
+        });
     }
-    labels
+
+    assert!(
+        all_fetched_hints.is_empty(),
+        "Did not expect background hints fetch tasks, but got {} of them",
+        all_fetched_hints.len()
+    );
+
+    all_cached_labels
 }
 
 #[track_caller]

crates/diagnostics/src/diagnostics_tests.rs

@@ -1,9 +1,9 @@
 use super::*;
 use collections::{HashMap, HashSet};
 use editor::{
-    DisplayPoint, EditorSettings,
+    DisplayPoint, EditorSettings, Inlay,
     actions::{GoToDiagnostic, GoToPreviousDiagnostic, Hover, MoveToBeginning},
-    display_map::{DisplayRow, Inlay},
+    display_map::DisplayRow,
     test::{
         editor_content_with_blocks, editor_lsp_test_context::EditorLspTestContext,
         editor_test_context::EditorTestContext,