Refresh code lens after dynamic registration

crates/editor/src/code_lens.rs

@@ -769,6 +769,173 @@ mod tests {
         });
     }
 
+    #[gpui::test]
+    async fn test_code_lens_refresh_requeries_open_document(cx: &mut TestAppContext) {
+        init_test(cx, |_| {});
+        update_test_editor_settings(cx, &|settings| {
+            settings.code_lens = Some(CodeLens::On);
+        });
+
+        let mut cx = EditorLspTestContext::new_typescript(
+            lsp::ServerCapabilities {
+                code_lens_provider: Some(lsp::CodeLensOptions {
+                    resolve_provider: None,
+                }),
+                execute_command_provider: Some(lsp::ExecuteCommandOptions {
+                    commands: vec!["lens_cmd".to_string()],
+                    ..lsp::ExecuteCommandOptions::default()
+                }),
+                ..lsp::ServerCapabilities::default()
+            },
+            cx,
+        )
+        .await;
+
+        let lens_title = Arc::new(Mutex::new("Initial lens".to_string()));
+        let mut code_lens_request =
+            cx.set_request_handler::<lsp::request::CodeLensRequest, _, _>({
+                let lens_title = lens_title.clone();
+                move |_, _, _| {
+                    let lens_title = lens_title.clone();
+                    async move {
+                        let title = lens_title.lock().unwrap().clone();
+                        Ok(Some(vec![lsp::CodeLens {
+                            range: lsp::Range::new(
+                                lsp::Position::new(0, 0),
+                                lsp::Position::new(0, 19),
+                            ),
+                            command: Some(lsp::Command {
+                                title,
+                                command: "lens_cmd".to_owned(),
+                                arguments: None,
+                            }),
+                            data: None,
+                        }]))
+                    }
+                }
+            });
+
+        cx.set_state("ˇfunction hello() {}\nfunction world() {}");
+        assert!(
+            code_lens_request.next().await.is_some(),
+            "should have received the initial code lens request"
+        );
+        cx.run_until_parked();
+        cx.editor(|editor, _, cx| {
+            assert_eq!(
+                code_lens_assertion_text(editor, cx),
+                indoc! {r#"
+                    Lenses: Initial lens
+                    Line 1: function hello() {}
+                "#},
+                "initial fetch should render the server title"
+            );
+        });
+
+        *lens_title.lock().unwrap() = "Refreshed lens".to_string();
+        cx.lsp
+            .request::<lsp::request::CodeLensRefresh>((), lsp::DEFAULT_LSP_REQUEST_TIMEOUT)
+            .await
+            .into_response()
+            .expect("code lens refresh request failed");
+        cx.executor()
+            .advance_clock(LSP_REQUEST_DEBOUNCE_TIMEOUT * 2);
+        cx.run_until_parked();
+        cx.editor(|editor, _, cx| {
+            assert_eq!(
+                code_lens_assertion_text(editor, cx),
+                indoc! {r#"
+                    Lenses: Refreshed lens
+                    Line 1: function hello() {}
+                "#},
+                "refresh should update the displayed lens to the new server title"
+            );
+        });
+    }
+
+    #[gpui::test]
+    async fn test_code_lens_dynamic_registration_requeries_open_document(cx: &mut TestAppContext) {
+        init_test(cx, |_| {});
+        update_test_editor_settings(cx, &|settings| {
+            settings.code_lens = Some(CodeLens::On);
+        });
+
+        // The server advertises no code lens capability up front; it registers
+        // `textDocument/codeLens` dynamically only after the document is open.
+        let mut cx = EditorLspTestContext::new_typescript(
+            lsp::ServerCapabilities {
+                execute_command_provider: Some(lsp::ExecuteCommandOptions {
+                    commands: vec!["lens_cmd".to_string()],
+                    ..lsp::ExecuteCommandOptions::default()
+                }),
+                ..lsp::ServerCapabilities::default()
+            },
+            cx,
+        )
+        .await;
+
+        let _code_lens_request =
+            cx.set_request_handler::<lsp::request::CodeLensRequest, _, _>(move |_, _, _| async {
+                Ok(Some(vec![lsp::CodeLens {
+                    range: lsp::Range::new(lsp::Position::new(0, 0), lsp::Position::new(0, 19)),
+                    command: Some(lsp::Command {
+                        title: "Dynamic lens".to_owned(),
+                        command: "lens_cmd".to_owned(),
+                        arguments: None,
+                    }),
+                    data: None,
+                }]))
+            });
+
+        cx.set_state("ˇfunction hello() {}\nfunction world() {}");
+        // Drain any debounced refresh scheduled before the capability exists, so
+        // the post-registration re-query can only come from the dynamic
+        // registration handling itself.
+        cx.executor()
+            .advance_clock(LSP_REQUEST_DEBOUNCE_TIMEOUT * 2);
+        cx.run_until_parked();
+        cx.editor(|editor, _, cx| {
+            assert_eq!(
+                code_lens_assertion_text(editor, cx),
+                "\n",
+                "no lenses should render before the capability is registered"
+            );
+        });
+
+        cx.lsp
+            .request::<lsp::request::RegisterCapability>(
+                lsp::RegistrationParams {
+                    registrations: vec![lsp::Registration {
+                        id: "code-lens".to_string(),
+                        method: "textDocument/codeLens".to_string(),
+                        register_options: Some(
+                            serde_json::to_value(lsp::CodeLensOptions {
+                                resolve_provider: None,
+                            })
+                            .unwrap(),
+                        ),
+                    }],
+                },
+                lsp::DEFAULT_LSP_REQUEST_TIMEOUT,
+            )
+            .await
+            .into_response()
+            .expect("register capability request failed");
+        cx.executor()
+            .advance_clock(LSP_REQUEST_DEBOUNCE_TIMEOUT * 2);
+        cx.run_until_parked();
+        cx.editor(|editor, _, cx| {
+            assert_eq!(
+                code_lens_assertion_text(editor, cx),
+                indoc! {r#"
+                    Lenses: Dynamic lens
+                    Line 1: function hello() {}
+                "#},
+                "dynamic textDocument/codeLens registration should re-query and display lenses for the open document"
+            );
+        });
+    }
+
     #[gpui::test]
     async fn test_code_lens_blocks_kept_across_refresh(cx: &mut TestAppContext) {
         init_test(cx, |_| {});

crates/project/src/lsp_store.rs

@@ -1134,23 +1134,12 @@ impl LocalLspStore {
 
         language_server
             .on_request::<lsp::request::CodeLensRefresh, _, _>({
-                let this = lsp_store.clone();
+                let lsp_store = lsp_store.clone();
                 move |(), cx| {
-                    let this = this.clone();
-                    let mut cx = cx.clone();
-                    async move {
-                        this.update(&mut cx, |this, cx| {
-                            this.invalidate_code_lens();
-                            cx.emit(LspStoreEvent::RefreshCodeLens);
-                            this.downstream_client.as_ref().map(|(client, project_id)| {
-                                client.send(proto::RefreshCodeLens {
-                                    project_id: *project_id,
-                                })
-                            })
-                        })?
-                        .transpose()?;
-                        Ok(())
-                    }
+                    let result = lsp_store.update(cx, |lsp_store, cx| {
+                        lsp_store.refresh_code_lens(cx);
+                    });
+                    async move { result }
                 }
             })
             .detach();
@@ -13128,6 +13117,7 @@ impl LspStore {
                             capabilities.code_lens_provider = Some(caps);
                         });
                         notify_server_capabilities_updated(&server, cx);
+                        self.refresh_code_lens(cx);
                     }
                 }
                 "textDocument/diagnostic" => {

crates/project/src/lsp_store/code_lens.rs

@@ -15,6 +15,7 @@ use rpc::{TypedEnvelope, proto};
 use settings::Settings as _;
 use std::time::Duration;
 use text::OffsetRangeExt as _;
+use util::ResultExt as _;
 
 use crate::{
     CodeAction, LspAction, LspStore, LspStoreEvent, Project,
@@ -71,10 +72,20 @@ fn flatten_cache(lens: &HashMap<LanguageServerId, CodeLensActions>) -> CodeLensA
 }
 
 impl LspStore {
-    pub(super) fn invalidate_code_lens(&mut self) {
+    pub(super) fn refresh_code_lens(&mut self, cx: &mut Context<Self>) {
         for lsp_data in self.lsp_data.values_mut() {
             lsp_data.code_lens = None;
         }
+
+        cx.emit(LspStoreEvent::RefreshCodeLens);
+        if let Some((downstream_client, project_id)) = self.downstream_client.as_ref() {
+            downstream_client
+                .send(proto::RefreshCodeLens {
+                    project_id: *project_id,
+                })
+                .context("sending refresh code lens downstream")
+                .log_err();
+        }
     }
 
     /// Fetches all code lenses for the buffer, each tagged with the
@@ -403,8 +414,7 @@ impl LspStore {
         mut cx: AsyncApp,
     ) -> Result<proto::Ack> {
         lsp_store.update(&mut cx, |lsp_store, cx| {
-            lsp_store.invalidate_code_lens();
-            cx.emit(LspStoreEvent::RefreshCodeLens);
+            lsp_store.refresh_code_lens(cx);
         });
         Ok(proto::Ack {})
     }