Fix code lens not being resolved in remote workflows

crates/collab/src/rpc.rs

@@ -363,6 +363,7 @@ impl Server {
             .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::ResolveInlayHint>)
+            .add_request_handler(forward_read_only_project_request::<proto::ResolveCodeAction>)
             .add_request_handler(forward_read_only_project_request::<proto::ResolveDocumentLink>)
             .add_request_handler(forward_read_only_project_request::<proto::GetColorPresentation>)
             .add_request_handler(forward_read_only_project_request::<proto::OpenBufferByPath>)

crates/collab/tests/integration/editor_tests.rs

@@ -10,6 +10,7 @@ use editor::{
         CopyFileName, CopyFileNameWithoutExtension, ExpandMacroRecursively, MoveToEnd, Redo,
         Rename, SelectAll, ToggleCodeActions, Undo,
     },
+    code_context_menus::CodeContextMenu,
     test::{
         editor_test_context::{AssertionContextManager, EditorTestContext},
         expand_macro_recursively,
@@ -1411,6 +1412,142 @@ async fn test_slow_lsp_server(cx_a: &mut TestAppContext, cx_b: &mut TestAppConte
     )
 }
 
+#[gpui::test]
+async fn test_collaborating_with_code_lens_resolve(
+    cx_a: &mut TestAppContext,
+    cx_b: &mut TestAppContext,
+) {
+    let mut server = TestServer::start(cx_a.executor()).await;
+    let client_a = server.create_client(cx_a, "user_a").await;
+    let client_b = server.create_client(cx_b, "user_b").await;
+    server
+        .create_room(&mut [(&client_a, cx_a), (&client_b, cx_b)])
+        .await;
+    let active_call_a = cx_a.read(ActiveCall::global);
+    cx_b.update(editor::init);
+    cx_b.update(|cx| {
+        SettingsStore::update_global(cx, |store, cx| {
+            store.update_user_settings(cx, |settings| {
+                settings.editor.code_lens = Some(settings::CodeLens::Menu);
+            });
+        });
+    });
+
+    let capabilities = lsp::ServerCapabilities {
+        code_lens_provider: Some(lsp::CodeLensOptions {
+            resolve_provider: Some(true),
+        }),
+        ..lsp::ServerCapabilities::default()
+    };
+    client_a.language_registry().add(rust_lang());
+    client_a.language_registry().register_fake_lsp(
+        "Rust",
+        FakeLspAdapter {
+            capabilities: capabilities.clone(),
+            initializer: Some(Box::new(|fake_lsp| {
+                fake_lsp.set_request_handler::<lsp::request::CodeLensRequest, _, _>(
+                    |_, _| async move {
+                        Ok(Some(vec![lsp::CodeLens {
+                            range: lsp::Range::new(
+                                lsp::Position::new(0, 0),
+                                lsp::Position::new(0, 9),
+                            ),
+                            command: None,
+                            data: Some(serde_json::json!({ "id": "lens" })),
+                        }]))
+                    },
+                );
+                fake_lsp.set_request_handler::<lsp::request::CodeLensResolve, _, _>(
+                    |lens, _| async move {
+                        Ok(lsp::CodeLens {
+                            command: Some(lsp::Command {
+                                title: "1 reference".to_string(),
+                                command: "noop".to_string(),
+                                arguments: None,
+                            }),
+                            ..lens
+                        })
+                    },
+                );
+            })),
+            ..FakeLspAdapter::default()
+        },
+    );
+    client_b.language_registry().add(rust_lang());
+    client_b.language_registry().register_fake_lsp_adapter(
+        "Rust",
+        FakeLspAdapter {
+            capabilities,
+            ..FakeLspAdapter::default()
+        },
+    );
+
+    client_a
+        .fs()
+        .insert_tree(
+            path!("/dir"),
+            json!({
+                "one.rs": "const ONE: usize = 1;"
+            }),
+        )
+        .await;
+    let (project_a, worktree_id) = client_a.build_local_project(path!("/dir"), cx_a).await;
+    let project_id = active_call_a
+        .update(cx_a, |call, cx| call.share_project(project_a.clone(), cx))
+        .await
+        .unwrap();
+    let project_b = client_b.join_remote_project(project_id, cx_b).await;
+
+    let (workspace_b, cx_b) = client_b.build_workspace(&project_b, cx_b);
+    let editor_b = workspace_b
+        .update_in(cx_b, |workspace, window, cx| {
+            workspace.open_path((worktree_id, rel_path("one.rs")), None, true, window, cx)
+        })
+        .await
+        .unwrap()
+        .downcast::<Editor>()
+        .unwrap();
+    cx_a.run_until_parked();
+    cx_b.run_until_parked();
+
+    editor_b.update_in(cx_b, |editor, window, cx| {
+        editor.change_selections(SelectionEffects::no_scroll(), window, cx, |s| {
+            s.select_ranges([Point::new(0, 0)..Point::new(0, 0)]);
+        });
+    });
+    cx_a.background_executor
+        .advance_clock(editor::CODE_ACTIONS_DEBOUNCE_TIMEOUT * 2);
+    cx_a.run_until_parked();
+    cx_b.run_until_parked();
+
+    editor_b.update_in(cx_b, |editor, window, cx| {
+        editor.toggle_code_actions(
+            &ToggleCodeActions {
+                deployed_from: None,
+                quick_launch: false,
+            },
+            window,
+            cx,
+        );
+    });
+    cx_a.run_until_parked();
+    cx_b.run_until_parked();
+
+    editor_b.update(cx_b, |editor, _| {
+        assert!(editor.context_menu_visible());
+        let menu = editor.context_menu().borrow();
+        let actions_menu = match menu.as_ref() {
+            Some(CodeContextMenu::CodeActions(m)) => m,
+            _ => panic!("Expected code actions menu to be visible"),
+        };
+        let item = actions_menu
+            .actions
+            .get(0)
+            .expect("Expected at least one item in menu");
+        assert_eq!(item.label(), "1 reference");
+    });
+}
+
 #[gpui::test(iterations = 10)]
 async fn test_language_server_statuses(cx_a: &mut TestAppContext, cx_b: &mut TestAppContext) {
     let mut server = TestServer::start(cx_a.executor()).await;

crates/project/src/lsp_store.rs

@@ -4218,6 +4218,7 @@ impl LspStore {
         client.add_entity_request_handler(Self::handle_apply_code_action);
         client.add_entity_request_handler(Self::handle_get_project_symbols);
         client.add_entity_request_handler(Self::handle_resolve_inlay_hint);
+        client.add_entity_request_handler(Self::handle_resolve_code_action);
         client.add_entity_request_handler(Self::handle_resolve_document_link);
         client.add_entity_request_handler(Self::handle_get_color_presentation);
         client.add_entity_request_handler(Self::handle_open_buffer_for_symbol);
@@ -5710,6 +5711,73 @@ impl LspStore {
         }
     }
 
+    pub fn resolve_code_action(
+        &self,
+        buffer: &Entity<Buffer>,
+        mut action: CodeAction,
+        cx: &mut Context<Self>,
+    ) -> Task<Result<CodeAction>> {
+        if action.resolved {
+            return Task::ready(Ok(action));
+        }
+        if let Some((upstream_client, project_id)) = self.upstream_client() {
+            let request = proto::ResolveCodeAction {
+                project_id,
+                buffer_id: buffer.read(cx).remote_id().into(),
+                action: Some(Self::serialize_code_action(&action)),
+            };
+            cx.background_spawn(async move {
+                let response = upstream_client
+                    .request(request)
+                    .await
+                    .context("resolve code action proto request")?;
+                let action = response.action.context("missing resolved action")?;
+                Self::deserialize_code_action(action)
+            })
+        } else if self.mode.is_local() {
+            let server_id = action.server_id;
+            let Some(lang_server) = buffer.update(cx, |buffer, cx| {
+                self.language_server_for_local_buffer(buffer, server_id, cx)
+                    .map(|(_, server)| server.clone())
+            }) else {
+                return Task::ready(Ok(action));
+            };
+            let request_timeout = ProjectSettings::get_global(cx)
+                .global_lsp_settings
+                .get_request_timeout();
+            cx.background_spawn(async move {
+                LocalLspStore::try_resolve_code_action(&lang_server, &mut action, request_timeout)
+                    .await
+                    .context("resolving a code action")?;
+                Ok(action)
+            })
+        } else {
+            Task::ready(Err(anyhow!("no upstream client and not local")))
+        }
+    }
+
+    pub(super) async fn handle_resolve_code_action(
+        lsp_store: Entity<Self>,
+        envelope: TypedEnvelope<proto::ResolveCodeAction>,
+        mut cx: AsyncApp,
+    ) -> Result<proto::ResolveCodeActionResponse> {
+        let action =
+            Self::deserialize_code_action(envelope.payload.action.context("invalid action")?)?;
+        let buffer = lsp_store.update(&mut cx, |lsp_store, cx| {
+            let buffer_id = BufferId::new(envelope.payload.buffer_id)?;
+            lsp_store.buffer_store.read(cx).get_existing(buffer_id)
+        })?;
+        let resolved = lsp_store
+            .update(&mut cx, |lsp_store, cx| {
+                lsp_store.resolve_code_action(&buffer, action, cx)
+            })
+            .await
+            .context("resolving code action")?;
+        Ok(proto::ResolveCodeActionResponse {
+            action: Some(Self::serialize_code_action(&resolved)),
+        })
+    }
+
     pub fn apply_code_action_kind(
         &mut self,
         buffers: HashSet<Entity<Buffer>>,

crates/project/src/lsp_store/code_lens.rs

@@ -3,7 +3,7 @@ use std::sync::Arc;
 
 use anyhow::{Context as _, Result};
 use clock::Global;
-use collections::HashMap;
+use collections::{HashMap, HashSet};
 use futures::{
     FutureExt as _,
     future::{Shared, join_all},
@@ -122,21 +122,23 @@ impl LspStore {
     ) -> CodeLensTask {
         let version_queried_for = buffer.read(cx).version();
         let buffer_id = buffer.read(cx).remote_id();
-        let existing_servers = self.as_local().map(|local| {
+        let existing_servers = if let Some(local) = self.as_local() {
             local
                 .buffers_opened_in_servers
                 .get(&buffer_id)
                 .cloned()
                 .unwrap_or_default()
-        });
+        } else {
+            self.relevant_server_ids_for_capability_check(buffer, cx)
+                .into_iter()
+                .collect()
+        };
 
         if let Some(lsp_data) = self.current_lsp_data(buffer_id) {
             if let Some(cached_lens) = &lsp_data.code_lens {
                 if !version_queried_for.changed_since(&lsp_data.buffer_version) {
-                    let has_different_servers = existing_servers.is_some_and(|existing_servers| {
-                        existing_servers != cached_lens.lens.keys().copied().collect()
-                    });
-                    if !has_different_servers {
+                    let cached_servers = cached_lens.lens.keys().copied().collect::<HashSet<_>>();
+                    if existing_servers == cached_servers {
                         return Task::ready(Ok(Some(flatten_cache(&cached_lens.lens)))).shared();
                     }
                 } else if let Some((updating_for, running_update)) = cached_lens.update.as_ref() {
@@ -304,7 +306,7 @@ impl LspStore {
     /// `(id, resolved_action)` pair is returned.
     ///
     /// All visibility / batching policy lives in the caller. Remote (proto)
-    /// resolves are not yet supported and currently yield `None`.
+    /// resolves are forwarded to the host via [`Self::resolve_code_action`].
     pub fn resolve_code_lens(
         &mut self,
         buffer: &Entity<Buffer>,
@@ -339,6 +341,47 @@ impl LspStore {
             return Task::ready(None).shared();
         };
         let lens = lens.clone();
+        let action = cached.clone();
+
+        if self.upstream_client().is_some() {
+            if !self.check_if_capable_for_proto_request(buffer, GetCodeLens::can_resolve_lens, cx) {
+                return Task::ready(None).shared();
+            }
+            let resolve = self.resolve_code_action(buffer, action, cx);
+            let task = cx
+                .spawn(async move |lsp_store, cx| {
+                    let resolved = resolve
+                        .await
+                        .context("resolving remote code lens")
+                        .log_err()?;
+                    lsp_store
+                        .update(cx, |lsp_store, _| {
+                            let code_lens = lsp_store
+                                .lsp_data
+                                .get_mut(&buffer_id)
+                                .and_then(|data| data.code_lens.as_mut())?;
+                            code_lens.resolving.remove(&key);
+                            let action = code_lens
+                                .lens
+                                .get_mut(&server_id)
+                                .and_then(|cache| cache.get_mut(&lens_id))?;
+                            action.resolved = true;
+                            action.lsp_action = resolved.lsp_action;
+                            Some((lens_id, action.clone()))
+                        })
+                        .ok()
+                        .flatten()
+                })
+                .shared();
+            if let Some(code_lens) = self
+                .lsp_data
+                .get_mut(&buffer_id)
+                .and_then(|data| data.code_lens.as_mut())
+            {
+                code_lens.resolving.insert(key, task.clone());
+            }
+            return task;
+        }
 
         let Some(server) = self.language_server_for_id(server_id) else {
             return Task::ready(None).shared();

crates/proto/proto/lsp.proto

@@ -378,6 +378,16 @@ message ApplyCodeAction {
   CodeAction action = 3;
 }
 
+message ResolveCodeAction {
+  uint64 project_id = 1;
+  uint64 buffer_id = 2;
+  CodeAction action = 3;
+}
+
+message ResolveCodeActionResponse {
+  CodeAction action = 1;
+}
+
 message ApplyCodeActionResponse {
   ProjectTransaction transaction = 1;
 }

crates/proto/proto/zed.proto

@@ -487,7 +487,9 @@ message Envelope {
     ResolveDocumentLinkResponse resolve_document_link_response = 456;
     GitWorktreeCreatedAt git_worktree_created_at = 457;
     GitWorktreeCreatedAtResponse git_worktree_created_at_response = 458;
-    TelemetryEvent telemetry_event = 459; // current max
+    TelemetryEvent telemetry_event = 459;
+    ResolveCodeAction resolve_code_action = 460;
+    ResolveCodeActionResponse resolve_code_action_response = 461; // current max
   }
 
   reserved 87 to 88;

crates/proto/src/proto.rs

@@ -233,6 +233,8 @@ messages!(
     (RefreshCodeLens, Background),
     (GetCodeLens, Background),
     (GetCodeLensResponse, Background),
+    (ResolveCodeAction, Background),
+    (ResolveCodeActionResponse, Background),
     (RespondToChannelInvite, Foreground),
     (RespondToContactRequest, Foreground),
     (RestartLanguageServers, Foreground),
@@ -472,6 +474,7 @@ request_messages!(
         ResolveCompletionDocumentationResponse
     ),
     (ResolveInlayHint, ResolveInlayHintResponse),
+    (ResolveCodeAction, ResolveCodeActionResponse),
     (GetDocumentColor, GetDocumentColorResponse),
     (GetDocumentLinks, GetDocumentLinksResponse),
     (ResolveDocumentLink, ResolveDocumentLinkResponse),
@@ -687,6 +690,7 @@ entity_messages!(
     RenameProjectEntry,
     ResolveCompletionDocumentation,
     ResolveInlayHint,
+    ResolveCodeAction,
     SaveBuffer,
     Stage,
     StartLanguageServer,