Handle dynamic registration of semantic tokens capability

crates/editor/src/semantic_tokens.rs

@@ -572,6 +572,107 @@ mod tests {
         assert_eq!(full_counter.load(atomic::Ordering::Acquire), 2);
     }
 
+    #[gpui::test]
+    async fn lsp_semantic_tokens_dynamic_registration_requeries_open_document(
+        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()
+                },
+            );
+        });
+
+        // The server advertises no semantic tokens capability up front; it only
+        // registers `textDocument/semanticTokens` dynamically, after the document
+        // is already open (as Roslyn does).
+        let mut cx = EditorLspTestContext::new_rust(lsp::ServerCapabilities::default(), cx).await;
+
+        let full_counter = Arc::new(AtomicUsize::new(0));
+        let _full_request = cx
+            .set_request_handler::<lsp::request::SemanticTokensFullRequest, _, _>({
+                let full_counter = full_counter.clone();
+                move |_, _, _| {
+                    full_counter.fetch_add(1, atomic::Ordering::Release);
+                    async move {
+                        Ok(Some(lsp::SemanticTokensResult::Tokens(
+                            lsp::SemanticTokens {
+                                data: vec![0, 3, 4, 0, 0],
+                                result_id: None,
+                            },
+                        )))
+                    }
+                }
+            });
+
+        cx.set_state("ˇfn main() {}");
+        // Drain the refresh scheduled on open (while no capability exists yet), so a
+        // later request can only come from the dynamic-registration refresh itself.
+        cx.executor().advance_clock(Duration::from_millis(200));
+        cx.run_until_parked();
+        assert_eq!(
+            full_counter.load(atomic::Ordering::Acquire),
+            0,
+            "no semantic tokens should be requested before the capability is registered"
+        );
+        assert!(
+            extract_semantic_highlights(&cx.editor, &cx).is_empty(),
+            "no semantic highlights before the capability is registered"
+        );
+
+        cx.lsp
+            .request::<lsp::request::RegisterCapability>(
+                lsp::RegistrationParams {
+                    registrations: vec![lsp::Registration {
+                        id: "semantic-tokens".to_string(),
+                        method: "textDocument/semanticTokens".to_string(),
+                        register_options: Some(
+                            serde_json::to_value(lsp::SemanticTokensRegistrationOptions {
+                                text_document_registration_options:
+                                    lsp::TextDocumentRegistrationOptions {
+                                        document_selector: None,
+                                    },
+                                semantic_tokens_options: lsp::SemanticTokensOptions {
+                                    legend: lsp::SemanticTokensLegend {
+                                        token_types: vec!["function".into()],
+                                        token_modifiers: Vec::new(),
+                                    },
+                                    full: Some(lsp::SemanticTokensFullOptions::Bool(true)),
+                                    ..lsp::SemanticTokensOptions::default()
+                                },
+                                static_registration_options: lsp::StaticRegistrationOptions {
+                                    id: None,
+                                },
+                            })
+                            .unwrap(),
+                        ),
+                    }],
+                },
+                lsp::DEFAULT_LSP_REQUEST_TIMEOUT,
+            )
+            .await
+            .into_response()
+            .expect("register capability request failed");
+
+        cx.executor().advance_clock(Duration::from_millis(200));
+        cx.run_until_parked();
+        assert!(
+            full_counter.load(atomic::Ordering::Acquire) >= 1,
+            "dynamic registration should re-query semantic tokens for the open document"
+        );
+
+        assert_eq!(
+            extract_semantic_highlights(&cx.editor, &cx),
+            vec![MultiBufferOffset(3)..MultiBufferOffset(7)],
+            "the open document should display semantic tokens after dynamic registration"
+        );
+    }
+
     #[gpui::test]
     async fn lsp_semantic_tokens_full_none_result_id(cx: &mut TestAppContext) {
         init_test(cx, |_| {});

crates/lsp/src/lsp.rs

@@ -934,7 +934,7 @@ impl LanguageServer {
                         dynamic_registration: Some(true),
                     }),
                     semantic_tokens: Some(SemanticTokensClientCapabilities {
-                        dynamic_registration: Some(false),
+                        dynamic_registration: Some(true),
                         requests: SemanticTokensClientCapabilitiesRequests {
                             range: None,
                             full: Some(SemanticTokensFullOptions::Delta { delta: Some(true) }),

crates/project/src/lsp_store.rs

@@ -1164,26 +1164,11 @@ impl LocalLspStore {
                     let request_id = request_id.clone();
                     let mut cx = cx.clone();
                     async move {
-                        lsp_store
-                            .update(&mut cx, |lsp_store, cx| {
-                                let request_id =
-                                    Some(request_id.fetch_add(1, atomic::Ordering::AcqRel));
-                                cx.emit(LspStoreEvent::RefreshSemanticTokens {
-                                    server_id,
-                                    request_id,
-                                });
-                                lsp_store
-                                    .downstream_client
-                                    .as_ref()
-                                    .map(|(client, project_id)| {
-                                        client.send(proto::RefreshSemanticTokens {
-                                            project_id: *project_id,
-                                            server_id: server_id.to_proto(),
-                                            request_id: request_id.map(|id| id as u64),
-                                        })
-                                    })
-                            })?
-                            .transpose()?;
+                        lsp_store.update(&mut cx, |lsp_store, cx| {
+                            let request_id =
+                                Some(request_id.fetch_add(1, atomic::Ordering::AcqRel));
+                            lsp_store.refresh_semantic_tokens(server_id, request_id, cx);
+                        })?;
                         Ok(())
                     }
                 }
@@ -13223,6 +13208,23 @@ impl LspStore {
                         notify_server_capabilities_updated(&server, cx);
                     }
                 }
+                "textDocument/semanticTokens" => {
+                    if let Some(caps) = reg
+                        .register_options
+                        .map(serde_json::from_value::<lsp::SemanticTokensRegistrationOptions>)
+                        .transpose()?
+                    {
+                        server.update_capabilities(|capabilities| {
+                            capabilities.semantic_tokens_provider = Some(
+                                lsp::SemanticTokensServerCapabilities::SemanticTokensRegistrationOptions(caps),
+                            );
+                        });
+                        notify_server_capabilities_updated(&server, cx);
+                        // Re-query already-open buffers, which would otherwise keep
+                        // tree-sitter-only highlighting until edited.
+                        self.refresh_semantic_tokens(server_id, None, cx);
+                    }
+                }
                 _ => log::warn!("unhandled capability registration: {reg:?}"),
             }
         }
@@ -13344,6 +13346,12 @@ impl LspStore {
                     });
                     notify_server_capabilities_updated(&server, cx);
                 }
+                "textDocument/semanticTokens" => {
+                    server.update_capabilities(|capabilities| {
+                        capabilities.semantic_tokens_provider = None;
+                    });
+                    notify_server_capabilities_updated(&server, cx);
+                }
                 "textDocument/didChange" => {
                     server.update_capabilities(|capabilities| {
                         let mut sync_options = Self::take_text_document_sync_options(capabilities);

crates/project/src/lsp_store/semantic_tokens.rs

@@ -351,6 +351,29 @@ impl LspStore {
         }
     }
 
+    /// `request_id` orders per-server refreshes (a higher id invalidates the cache).
+    /// Client-initiated refreshes (e.g. after dynamic registration) pass `None`.
+    pub(crate) fn refresh_semantic_tokens(
+        &mut self,
+        server_id: LanguageServerId,
+        request_id: Option<usize>,
+        cx: &mut Context<Self>,
+    ) {
+        cx.emit(LspStoreEvent::RefreshSemanticTokens {
+            server_id,
+            request_id,
+        });
+        if let Some((client, project_id)) = self.downstream_client.as_ref() {
+            client
+                .send(proto::RefreshSemanticTokens {
+                    project_id: *project_id,
+                    server_id: server_id.to_proto(),
+                    request_id: request_id.map(|id| id as u64),
+                })
+                .log_err();
+        }
+    }
+
     pub(crate) async fn handle_refresh_semantic_tokens(
         lsp_store: Entity<Self>,
         envelope: TypedEnvelope<proto::RefreshSemanticTokens>,

crates/project/tests/integration/project_tests.rs

@@ -2251,6 +2251,131 @@ async fn test_rescan_fs_change_is_reported_to_language_servers_as_changed(
     );
 }
 
+#[gpui::test]
+async fn test_dynamic_semantic_tokens_registration(cx: &mut gpui::TestAppContext) {
+    init_test(cx);
+
+    let fs = FakeFs::new(cx.executor());
+    fs.insert_tree(
+        path!("/the-root"),
+        json!({
+            "a.rs": "fn main() {}",
+        }),
+    )
+    .await;
+
+    let project = Project::test(fs.clone(), [path!("/the-root").as_ref()], cx).await;
+    let language_registry = project.read_with(cx, |project, _| project.languages().clone());
+    language_registry.add(rust_lang());
+    let mut fake_servers = language_registry.register_fake_lsp(
+        "Rust",
+        FakeLspAdapter {
+            name: "the-language-server",
+            // Crucially, no `semantic_tokens_provider` is advertised statically; the
+            // server only offers it through dynamic registration (as Roslyn does).
+            ..Default::default()
+        },
+    );
+
+    let _buffer = project
+        .update(cx, |project, cx| {
+            project.open_local_buffer_with_lsp(path!("/the-root/a.rs"), cx)
+        })
+        .await
+        .unwrap();
+
+    let fake_server = fake_servers.next().await.unwrap();
+    let server_id = fake_server.server.server_id();
+    cx.executor().run_until_parked();
+
+    let semantic_tokens_provider = |cx: &mut gpui::TestAppContext| {
+        project.read_with(cx, |project, cx| {
+            project
+                .lsp_store()
+                .read(cx)
+                .lsp_server_capabilities
+                .get(&server_id)
+                .and_then(|capabilities| capabilities.semantic_tokens_provider.clone())
+        })
+    };
+
+    assert!(
+        semantic_tokens_provider(cx).is_none(),
+        "server should not advertise semantic tokens before dynamic registration"
+    );
+
+    fake_server
+        .request::<lsp::request::RegisterCapability>(
+            lsp::RegistrationParams {
+                registrations: vec![lsp::Registration {
+                    id: "semantic-tokens".to_string(),
+                    method: "textDocument/semanticTokens".to_string(),
+                    register_options: serde_json::to_value(
+                        lsp::SemanticTokensRegistrationOptions {
+                            text_document_registration_options:
+                                lsp::TextDocumentRegistrationOptions {
+                                    document_selector: None,
+                                },
+                            semantic_tokens_options: lsp::SemanticTokensOptions {
+                                legend: lsp::SemanticTokensLegend {
+                                    token_types: vec!["keyword".into(), "variable".into()],
+                                    token_modifiers: vec![],
+                                },
+                                full: Some(lsp::SemanticTokensFullOptions::Bool(true)),
+                                ..Default::default()
+                            },
+                            static_registration_options: lsp::StaticRegistrationOptions {
+                                id: None,
+                            },
+                        },
+                    )
+                    .ok(),
+                }],
+            },
+            DEFAULT_LSP_REQUEST_TIMEOUT,
+        )
+        .await
+        .into_response()
+        .unwrap();
+    cx.executor().run_until_parked();
+
+    let provider = semantic_tokens_provider(cx)
+        .expect("semantic tokens provider should be set after dynamic registration");
+    // The capability round-trips through capability-sync serialization, which may
+    // normalize the registration options into plain options; either shape is fine
+    // as long as the legend survives.
+    let legend = match provider {
+        lsp::SemanticTokensServerCapabilities::SemanticTokensOptions(options) => options.legend,
+        lsp::SemanticTokensServerCapabilities::SemanticTokensRegistrationOptions(options) => {
+            options.semantic_tokens_options.legend
+        }
+    };
+    assert_eq!(
+        legend.token_types,
+        vec!["keyword".into(), "variable".into()],
+    );
+
+    fake_server
+        .request::<lsp::request::UnregisterCapability>(
+            lsp::UnregistrationParams {
+                unregisterations: vec![lsp::Unregistration {
+                    id: "semantic-tokens".to_string(),
+                    method: "textDocument/semanticTokens".to_string(),
+                }],
+            },
+            DEFAULT_LSP_REQUEST_TIMEOUT,
+        )
+        .await
+        .into_response()
+        .unwrap();
+    cx.executor().run_until_parked();
+
+    assert!(
+        semantic_tokens_provider(cx).is_none(),
+        "semantic tokens provider should be cleared after unregistration"
+    );
+}
+
 #[gpui::test]
 async fn test_reporting_fs_changes_to_language_servers(cx: &mut gpui::TestAppContext) {
     init_test(cx);