Properly discard tokens on language server stop (#48490)

Kirill Bulatov created 2 months ago

Follow-up of https://github.com/zed-industries/zed/pull/46356

Release Notes:

- N/A

Change summary

crates/editor/src/editor.rs          |   2 
crates/editor/src/semantic_tokens.rs | 111 +++++++++++++++++++++++++----
2 files changed, 97 insertions(+), 16 deletions(-)

Detailed changes

crates/editor/src/editor.rs

@@ -2163,7 +2163,7 @@ impl Editor {
                         }
                         editor.registered_buffers.clear();
                         editor.register_visible_buffers(cx);
-                        editor.update_semantic_tokens(None, None, cx);
+                        editor.refresh_semantic_token_highlights(cx);
                         editor.refresh_inlay_hints(InlayHintRefreshReason::ServerRemoved, cx);
                     }
                     project::Event::LanguageServerAdded(..) => {

crates/editor/src/semantic_tokens.rs

@@ -209,6 +209,9 @@ impl Editor {
                             tokens
                         },
                         Ok(BufferSemanticTokens { tokens: None }) => {
+                            editor.display_map.update(cx, |display_map, _| {
+                                display_map.invalidate_semantic_highlights(buffer_id);
+                            });
                             continue;
                         },
                         Err(e) => {
@@ -458,14 +461,14 @@ mod tests {
                         lsp::SemanticTokensOptions {
                             legend: lsp::SemanticTokensLegend {
                                 token_types: vec!["function".into()],
-                                token_modifiers: vec![],
+                                token_modifiers: Vec::new(),
                             },
                             full: Some(lsp::SemanticTokensFullOptions::Delta { delta: None }),
-                            ..Default::default()
+                            ..lsp::SemanticTokensOptions::default()
                         },
                     ),
                 ),
-                ..Default::default()
+                ..lsp::ServerCapabilities::default()
             },
             cx,
         )
@@ -536,14 +539,14 @@ mod tests {
                         lsp::SemanticTokensOptions {
                             legend: lsp::SemanticTokensLegend {
                                 token_types: vec!["function".into()],
-                                token_modifiers: vec![],
+                                token_modifiers: Vec::new(),
                             },
                             full: Some(lsp::SemanticTokensFullOptions::Delta { delta: Some(true) }),
-                            ..Default::default()
+                            ..lsp::SemanticTokensOptions::default()
                         },
                     ),
                 ),
-                ..Default::default()
+                ..lsp::ServerCapabilities::default()
             },
             cx,
         )
@@ -612,14 +615,14 @@ mod tests {
                         lsp::SemanticTokensOptions {
                             legend: lsp::SemanticTokensLegend {
                                 token_types: vec!["function".into()],
-                                token_modifiers: vec![],
+                                token_modifiers: Vec::new(),
                             },
                             full: Some(lsp::SemanticTokensFullOptions::Delta { delta: Some(true) }),
-                            ..Default::default()
+                            ..lsp::SemanticTokensOptions::default()
                         },
                     ),
                 ),
-                ..Default::default()
+                ..lsp::ServerCapabilities::default()
             },
             cx,
         )
@@ -659,7 +662,7 @@ mod tests {
                     async move {
                         Ok(Some(lsp::SemanticTokensFullDeltaResult::TokensDelta(
                             lsp::SemanticTokensDelta {
-                                edits: vec![],
+                                edits: Vec::new(),
                                 result_id: Some("b".into()),
                             },
                         )))
@@ -716,11 +719,11 @@ mod tests {
         // We have 2 language servers for TOML in this test.
         let toml_legend_1 = lsp::SemanticTokensLegend {
             token_types: vec!["property".into()],
-            token_modifiers: vec![],
+            token_modifiers: Vec::new(),
         };
         let toml_legend_2 = lsp::SemanticTokensLegend {
             token_types: vec!["number".into()],
-            token_modifiers: vec![],
+            token_modifiers: Vec::new(),
         };
 
         let app_state = cx.update(workspace::AppState::test);
@@ -954,11 +957,11 @@ mod tests {
 
         let toml_legend = lsp::SemanticTokensLegend {
             token_types: vec!["property".into()],
-            token_modifiers: vec![],
+            token_modifiers: Vec::new(),
         };
         let rust_legend = lsp::SemanticTokensLegend {
             token_types: vec!["constant".into()],
-            token_modifiers: vec![],
+            token_modifiers: Vec::new(),
         };
 
         let app_state = cx.update(workspace::AppState::test);
@@ -1221,7 +1224,7 @@ mod tests {
 
         let toml_legend = lsp::SemanticTokensLegend {
             token_types: vec!["property".into()],
-            token_modifiers: vec![],
+            token_modifiers: Vec::new(),
         };
 
         let app_state = cx.update(workspace::AppState::test);
@@ -1806,6 +1809,84 @@ mod tests {
         );
     }
 
+    #[gpui::test]
+    async fn test_stopping_language_server_clears_semantic_tokens(cx: &mut TestAppContext) {
+        init_test(cx, |_| {});
+
+        update_test_language_settings(cx, |language_settings| {
+            language_settings.languages.0.insert(
+                "Rust".into(),
+                LanguageSettingsContent {
+                    semantic_tokens: Some(SemanticTokens::Full),
+                    ..LanguageSettingsContent::default()
+                },
+            );
+        });
+
+        let mut cx = EditorLspTestContext::new_rust(
+            lsp::ServerCapabilities {
+                semantic_tokens_provider: Some(
+                    lsp::SemanticTokensServerCapabilities::SemanticTokensOptions(
+                        lsp::SemanticTokensOptions {
+                            legend: lsp::SemanticTokensLegend {
+                                token_types: vec!["function".into()],
+                                token_modifiers: Vec::new(),
+                            },
+                            full: Some(lsp::SemanticTokensFullOptions::Delta { delta: None }),
+                            ..lsp::SemanticTokensOptions::default()
+                        },
+                    ),
+                ),
+                ..lsp::ServerCapabilities::default()
+            },
+            cx,
+        )
+        .await;
+
+        let mut full_request = cx
+            .set_request_handler::<lsp::request::SemanticTokensFullRequest, _, _>(
+                move |_, _, _| async move {
+                    Ok(Some(lsp::SemanticTokensResult::Tokens(
+                        lsp::SemanticTokens {
+                            data: vec![
+                                0, // delta_line
+                                3, // delta_start
+                                4, // length
+                                0, // token_type
+                                0, // token_modifiers_bitset
+                            ],
+                            result_id: None,
+                        },
+                    )))
+                },
+            );
+
+        cx.set_state("ˇfn main() {}");
+        assert!(full_request.next().await.is_some());
+        cx.run_until_parked();
+
+        assert_eq!(
+            extract_semantic_highlights(&cx.editor, &cx),
+            vec![MultiBufferOffset(3)..MultiBufferOffset(7)],
+            "Semantic tokens should be present before stopping the server"
+        );
+
+        cx.update_editor(|editor, _, cx| {
+            let buffers = editor.buffer.read(cx).all_buffers().into_iter().collect();
+            editor.project.as_ref().unwrap().update(cx, |project, cx| {
+                project.stop_language_servers_for_buffers(buffers, HashSet::default(), cx);
+            })
+        });
+        cx.executor().advance_clock(Duration::from_millis(200));
+        cx.run_until_parked();
+
+        assert_eq!(
+            extract_semantic_highlights(&cx.editor, &cx),
+            Vec::new(),
+            "Semantic tokens should be cleared after stopping the server"
+        );
+    }
+
     fn extract_semantic_highlight_styles(
         editor: &Entity<Editor>,
         cx: &TestAppContext,