editor: Add Organize Imports Action (#25793)

Closes #10004

This PR adds support for the organize imports action. Previously, you
had to manually configure it in the settings and then use format to run
it.

Note: Default key binding will be `alt-shift-o` which is similar to
VSCode's organize import. Also, because `cmd-shift-o` is taken by
outline picker.

Todo:

- [x] Initial working
- [x] Handle remote
- [x] Handle multi buffer
- [x] Can we make it generic for executing any code action?

Release Notes:

- Added `editor:OrganizeImports` action to organize imports (sort,
remove unused, etc) for supported LSPs. You can trigger it by using the
`alt-shift-o` key binding.

assets/keymaps/default-macos.json

@@ -118,6 +118,7 @@
       "cmd-a": "editor::SelectAll",
       "cmd-l": "editor::SelectLine",
       "cmd-shift-i": "editor::Format",
+      "alt-shift-o": "editor::OrganizeImports",
       "cmd-shift-left": ["editor::SelectToBeginningOfLine", { "stop_at_soft_wraps": true, "stop_at_indent": true }],
       "shift-home": ["editor::SelectToBeginningOfLine", { "stop_at_soft_wraps": true, "stop_at_indent": true }],
       "ctrl-shift-a": ["editor::SelectToBeginningOfLine", { "stop_at_soft_wraps": true, "stop_at_indent": true }],

crates/collab/src/rpc.rs

@@ -328,6 +328,7 @@ impl Server {
             .add_request_handler(forward_mutating_project_request::<proto::PrepareRename>)
             .add_request_handler(forward_mutating_project_request::<proto::PerformRename>)
             .add_request_handler(forward_mutating_project_request::<proto::ReloadBuffers>)
+            .add_request_handler(forward_mutating_project_request::<proto::ApplyCodeActionKind>)
             .add_request_handler(forward_mutating_project_request::<proto::FormatBuffers>)
             .add_request_handler(forward_mutating_project_request::<proto::CreateProjectEntry>)
             .add_request_handler(forward_mutating_project_request::<proto::RenameProjectEntry>)

crates/editor/src/actions.rs

@@ -348,6 +348,7 @@ gpui::actions!(
         OpenPermalinkToLine,
         OpenSelectionsInMultibuffer,
         OpenUrl,
+        OrganizeImports,
         Outdent,
         AutoIndent,
         PageDown,

crates/editor/src/editor.rs

@@ -120,8 +120,8 @@ use task::{ResolvedTask, TaskTemplate, TaskVariables};
 use hover_links::{find_file, HoverLink, HoveredLinkState, InlayHighlight};
 pub use lsp::CompletionContext;
 use lsp::{
-    CompletionItemKind, CompletionTriggerKind, DiagnosticSeverity, InsertTextFormat,
+    CodeActionKind, CompletionItemKind, CompletionTriggerKind, DiagnosticSeverity,
-    LanguageServerId, LanguageServerName,
+    InsertTextFormat, LanguageServerId, LanguageServerName,
 };
 
 use language::BufferSnapshot;
@@ -203,6 +203,7 @@ pub(crate) const CURSORS_VISIBLE_FOR: Duration = Duration::from_millis(2000);
 #[doc(hidden)]
 pub const CODE_ACTIONS_DEBOUNCE_TIMEOUT: Duration = Duration::from_millis(250);
 
+pub(crate) const CODE_ACTION_TIMEOUT: Duration = Duration::from_secs(5);
 pub(crate) const FORMAT_TIMEOUT: Duration = Duration::from_secs(5);
 pub(crate) const SCROLL_CENTER_TOP_BOTTOM_DEBOUNCE_TIMEOUT: Duration = Duration::from_secs(1);
 
@@ -12494,7 +12495,6 @@ impl Editor {
                             buffer.push_transaction(&transaction.0, cx);
                         }
                     }
-
                     cx.notify();
                 })
                 .ok();
@@ -12503,6 +12503,60 @@ impl Editor {
         })
     }
 
+    fn organize_imports(
+        &mut self,
+        _: &OrganizeImports,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) -> Option<Task<Result<()>>> {
+        let project = match &self.project {
+            Some(project) => project.clone(),
+            None => return None,
+        };
+        Some(self.perform_code_action_kind(
+            project,
+            CodeActionKind::SOURCE_ORGANIZE_IMPORTS,
+            window,
+            cx,
+        ))
+    }
+
+    fn perform_code_action_kind(
+        &mut self,
+        project: Entity<Project>,
+        kind: CodeActionKind,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) -> Task<Result<()>> {
+        let buffer = self.buffer.clone();
+        let buffers = buffer.read(cx).all_buffers();
+        let mut timeout = cx.background_executor().timer(CODE_ACTION_TIMEOUT).fuse();
+        let apply_action = project.update(cx, |project, cx| {
+            project.apply_code_action_kind(buffers, kind, true, cx)
+        });
+        cx.spawn_in(window, |_, mut cx| async move {
+            let transaction = futures::select_biased! {
+                () = timeout => {
+                    log::warn!("timed out waiting for executing code action");
+                    None
+                }
+                transaction = apply_action.log_err().fuse() => transaction,
+            };
+            buffer
+                .update(&mut cx, |buffer, cx| {
+                    // check if we need this
+                    if let Some(transaction) = transaction {
+                        if !buffer.is_singleton() {
+                            buffer.push_transaction(&transaction.0, cx);
+                        }
+                    }
+                    cx.notify();
+                })
+                .ok();
+            Ok(())
+        })
+    }
+
     fn restart_language_server(
         &mut self,
         _: &RestartLanguageServer,

crates/editor/src/editor_tests.rs

@@ -7875,6 +7875,157 @@ async fn test_document_format_manual_trigger(cx: &mut TestAppContext) {
     );
 }
 
+#[gpui::test]
+async fn test_organize_imports_manual_trigger(cx: &mut TestAppContext) {
+    init_test(cx, |settings| {
+        settings.defaults.formatter = Some(language_settings::SelectedFormatter::List(
+            FormatterList(vec![Formatter::LanguageServer { name: None }].into()),
+        ))
+    });
+
+    let fs = FakeFs::new(cx.executor());
+    fs.insert_file(path!("/file.ts"), Default::default()).await;
+
+    let project = Project::test(fs, [path!("/").as_ref()], cx).await;
+
+    let language_registry = project.read_with(cx, |project, _| project.languages().clone());
+    language_registry.add(Arc::new(Language::new(
+        LanguageConfig {
+            name: "TypeScript".into(),
+            matcher: LanguageMatcher {
+                path_suffixes: vec!["ts".to_string()],
+                ..Default::default()
+            },
+            ..LanguageConfig::default()
+        },
+        Some(tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into()),
+    )));
+    update_test_language_settings(cx, |settings| {
+        settings.defaults.prettier = Some(PrettierSettings {
+            allowed: true,
+            ..PrettierSettings::default()
+        });
+    });
+    let mut fake_servers = language_registry.register_fake_lsp(
+        "TypeScript",
+        FakeLspAdapter {
+            capabilities: lsp::ServerCapabilities {
+                code_action_provider: Some(lsp::CodeActionProviderCapability::Simple(true)),
+                ..Default::default()
+            },
+            ..Default::default()
+        },
+    );
+
+    let buffer = project
+        .update(cx, |project, cx| {
+            project.open_local_buffer(path!("/file.ts"), cx)
+        })
+        .await
+        .unwrap();
+
+    let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
+    let (editor, cx) = cx.add_window_view(|window, cx| {
+        build_editor_with_project(project.clone(), buffer, window, cx)
+    });
+    editor.update_in(cx, |editor, window, cx| {
+        editor.set_text(
+            "import { a } from 'module';\nimport { b } from 'module';\n\nconst x = a;\n",
+            window,
+            cx,
+        )
+    });
+
+    cx.executor().start_waiting();
+    let fake_server = fake_servers.next().await.unwrap();
+
+    let format = editor
+        .update_in(cx, |editor, window, cx| {
+            editor.perform_code_action_kind(
+                project.clone(),
+                CodeActionKind::SOURCE_ORGANIZE_IMPORTS,
+                window,
+                cx,
+            )
+        })
+        .unwrap();
+    fake_server
+        .handle_request::<lsp::request::CodeActionRequest, _, _>(move |params, _| async move {
+            assert_eq!(
+                params.text_document.uri,
+                lsp::Url::from_file_path(path!("/file.ts")).unwrap()
+            );
+            Ok(Some(vec![lsp::CodeActionOrCommand::CodeAction(
+                lsp::CodeAction {
+                    title: "Organize Imports".to_string(),
+                    kind: Some(lsp::CodeActionKind::SOURCE_ORGANIZE_IMPORTS),
+                    edit: Some(lsp::WorkspaceEdit {
+                        changes: Some(
+                            [(
+                                params.text_document.uri.clone(),
+                                vec![lsp::TextEdit::new(
+                                    lsp::Range::new(
+                                        lsp::Position::new(1, 0),
+                                        lsp::Position::new(2, 0),
+                                    ),
+                                    "".to_string(),
+                                )],
+                            )]
+                            .into_iter()
+                            .collect(),
+                        ),
+                        ..Default::default()
+                    }),
+                    ..Default::default()
+                },
+            )]))
+        })
+        .next()
+        .await;
+    cx.executor().start_waiting();
+    format.await;
+    assert_eq!(
+        editor.update(cx, |editor, cx| editor.text(cx)),
+        "import { a } from 'module';\n\nconst x = a;\n"
+    );
+
+    editor.update_in(cx, |editor, window, cx| {
+        editor.set_text(
+            "import { a } from 'module';\nimport { b } from 'module';\n\nconst x = a;\n",
+            window,
+            cx,
+        )
+    });
+    // Ensure we don't lock if code action hangs.
+    fake_server.handle_request::<lsp::request::CodeActionRequest, _, _>(
+        move |params, _| async move {
+            assert_eq!(
+                params.text_document.uri,
+                lsp::Url::from_file_path(path!("/file.ts")).unwrap()
+            );
+            futures::future::pending::<()>().await;
+            unreachable!()
+        },
+    );
+    let format = editor
+        .update_in(cx, |editor, window, cx| {
+            editor.perform_code_action_kind(
+                project,
+                CodeActionKind::SOURCE_ORGANIZE_IMPORTS,
+                window,
+                cx,
+            )
+        })
+        .unwrap();
+    cx.executor().advance_clock(super::CODE_ACTION_TIMEOUT);
+    cx.executor().start_waiting();
+    format.await;
+    assert_eq!(
+        editor.update(cx, |editor, cx| editor.text(cx)),
+        "import { a } from 'module';\nimport { b } from 'module';\n\nconst x = a;\n"
+    );
+}
+
 #[gpui::test]
 async fn test_concurrent_format_requests(cx: &mut TestAppContext) {
     init_test(cx, |_| {});

crates/editor/src/element.rs

@@ -429,6 +429,13 @@ impl EditorElement {
                 cx.propagate();
             }
         });
+        register_action(editor, window, |editor, action, window, cx| {
+            if let Some(task) = editor.organize_imports(action, window, cx) {
+                task.detach_and_notify_err(window, cx);
+            } else {
+                cx.propagate();
+            }
+        });
         register_action(editor, window, Editor::restart_language_server);
         register_action(editor, window, Editor::show_character_palette);
         register_action(editor, window, |editor, action, window, cx| {

crates/project/src/lsp_store.rs

@@ -1089,6 +1089,64 @@ impl LocalLspStore {
         self.language_servers_for_buffer(buffer, cx).next()
     }
 
+    async fn execute_code_action_kind_locally(
+        lsp_store: WeakEntity<LspStore>,
+        mut buffers: Vec<Entity<Buffer>>,
+        kind: CodeActionKind,
+        push_to_history: bool,
+        mut cx: AsyncApp,
+    ) -> anyhow::Result<ProjectTransaction> {
+        // Do not allow multiple concurrent code actions requests for the
+        // same buffer.
+        lsp_store.update(&mut cx, |this, cx| {
+            let this = this.as_local_mut().unwrap();
+            buffers.retain(|buffer| {
+                this.buffers_being_formatted
+                    .insert(buffer.read(cx).remote_id())
+            });
+        })?;
+        let _cleanup = defer({
+            let this = lsp_store.clone();
+            let mut cx = cx.clone();
+            let buffers = &buffers;
+            move || {
+                this.update(&mut cx, |this, cx| {
+                    let this = this.as_local_mut().unwrap();
+                    for buffer in buffers {
+                        this.buffers_being_formatted
+                            .remove(&buffer.read(cx).remote_id());
+                    }
+                })
+                .ok();
+            }
+        });
+        let mut project_transaction = ProjectTransaction::default();
+
+        for buffer in &buffers {
+            let adapters_and_servers = lsp_store.update(&mut cx, |lsp_store, cx| {
+                buffer.update(cx, |buffer, cx| {
+                    lsp_store
+                        .as_local()
+                        .unwrap()
+                        .language_servers_for_buffer(buffer, cx)
+                        .map(|(adapter, lsp)| (adapter.clone(), lsp.clone()))
+                        .collect::<Vec<_>>()
+                })
+            })?;
+            Self::execute_code_actions_on_servers(
+                &lsp_store,
+                &adapters_and_servers,
+                vec![kind.clone()],
+                &buffer,
+                push_to_history,
+                &mut project_transaction,
+                &mut cx,
+            )
+            .await?;
+        }
+        Ok(project_transaction)
+    }
+
     async fn format_locally(
         lsp_store: WeakEntity<LspStore>,
         mut buffers: Vec<FormattableBuffer>,
@@ -2900,6 +2958,7 @@ impl LspStore {
         client.add_entity_message_handler(Self::handle_language_server_log);
         client.add_entity_message_handler(Self::handle_update_diagnostic_summary);
         client.add_entity_request_handler(Self::handle_format_buffers);
+        client.add_entity_request_handler(Self::handle_apply_code_action_kind);
         client.add_entity_request_handler(Self::handle_resolve_completion_documentation);
         client.add_entity_request_handler(Self::handle_apply_code_action);
         client.add_entity_request_handler(Self::handle_inlay_hints);
@@ -3891,6 +3950,65 @@ impl LspStore {
         }
     }
 
+    pub fn apply_code_action_kind(
+        &mut self,
+        buffers: HashSet<Entity<Buffer>>,
+        kind: CodeActionKind,
+        push_to_history: bool,
+        cx: &mut Context<Self>,
+    ) -> Task<anyhow::Result<ProjectTransaction>> {
+        if let Some(_) = self.as_local() {
+            cx.spawn(move |lsp_store, mut cx| async move {
+                let buffers = buffers.into_iter().collect::<Vec<_>>();
+                let result = LocalLspStore::execute_code_action_kind_locally(
+                    lsp_store.clone(),
+                    buffers,
+                    kind,
+                    push_to_history,
+                    cx.clone(),
+                )
+                .await;
+                lsp_store.update(&mut cx, |lsp_store, _| {
+                    lsp_store.update_last_formatting_failure(&result);
+                })?;
+                result
+            })
+        } else if let Some((client, project_id)) = self.upstream_client() {
+            let buffer_store = self.buffer_store();
+            cx.spawn(move |lsp_store, mut cx| async move {
+                let result = client
+                    .request(proto::ApplyCodeActionKind {
+                        project_id,
+                        kind: kind.as_str().to_owned(),
+                        buffer_ids: buffers
+                            .iter()
+                            .map(|buffer| {
+                                buffer.update(&mut cx, |buffer, _| buffer.remote_id().into())
+                            })
+                            .collect::<Result<_>>()?,
+                    })
+                    .await
+                    .and_then(|result| result.transaction.context("missing transaction"));
+                lsp_store.update(&mut cx, |lsp_store, _| {
+                    lsp_store.update_last_formatting_failure(&result);
+                })?;
+
+                let transaction_response = result?;
+                buffer_store
+                    .update(&mut cx, |buffer_store, cx| {
+                        buffer_store.deserialize_project_transaction(
+                            transaction_response,
+                            push_to_history,
+                            cx,
+                        )
+                    })?
+                    .await
+            })
+        } else {
+            Task::ready(Ok(ProjectTransaction::default()))
+        }
+    }
+
     pub fn resolve_inlay_hint(
         &self,
         hint: InlayHint,
@@ -7229,6 +7347,48 @@ impl LspStore {
         })
     }
 
+    async fn handle_apply_code_action_kind(
+        this: Entity<Self>,
+        envelope: TypedEnvelope<proto::ApplyCodeActionKind>,
+        mut cx: AsyncApp,
+    ) -> Result<proto::ApplyCodeActionKindResponse> {
+        let sender_id = envelope.original_sender_id().unwrap_or_default();
+        let format = this.update(&mut cx, |this, cx| {
+            let mut buffers = HashSet::default();
+            for buffer_id in &envelope.payload.buffer_ids {
+                let buffer_id = BufferId::new(*buffer_id)?;
+                buffers.insert(this.buffer_store.read(cx).get_existing(buffer_id)?);
+            }
+            let kind = match envelope.payload.kind.as_str() {
+                "" => Ok(CodeActionKind::EMPTY),
+                "quickfix" => Ok(CodeActionKind::QUICKFIX),
+                "refactor" => Ok(CodeActionKind::REFACTOR),
+                "refactor.extract" => Ok(CodeActionKind::REFACTOR_EXTRACT),
+                "refactor.inline" => Ok(CodeActionKind::REFACTOR_INLINE),
+                "refactor.rewrite" => Ok(CodeActionKind::REFACTOR_REWRITE),
+                "source" => Ok(CodeActionKind::SOURCE),
+                "source.organizeImports" => Ok(CodeActionKind::SOURCE_ORGANIZE_IMPORTS),
+                "source.fixAll" => Ok(CodeActionKind::SOURCE_FIX_ALL),
+                _ => Err(anyhow!("Invalid code action kind")),
+            }?;
+            anyhow::Ok(this.apply_code_action_kind(buffers, kind, false, cx))
+        })??;
+
+        let project_transaction = format.await?;
+        let project_transaction = this.update(&mut cx, |this, cx| {
+            this.buffer_store.update(cx, |buffer_store, cx| {
+                buffer_store.serialize_project_transaction_for_peer(
+                    project_transaction,
+                    sender_id,
+                    cx,
+                )
+            })
+        })?;
+        Ok(proto::ApplyCodeActionKindResponse {
+            transaction: Some(project_transaction),
+        })
+    }
+
     async fn shutdown_language_server(
         server_state: Option<LanguageServerState>,
         name: LanguageServerName,

crates/project/src/project.rs

@@ -3029,6 +3029,18 @@ impl Project {
         })
     }
 
+    pub fn apply_code_action_kind(
+        &self,
+        buffers: HashSet<Entity<Buffer>>,
+        kind: CodeActionKind,
+        push_to_history: bool,
+        cx: &mut Context<Self>,
+    ) -> Task<Result<ProjectTransaction>> {
+        self.lsp_store.update(cx, |lsp_store, cx| {
+            lsp_store.apply_code_action_kind(buffers, kind, push_to_history, cx)
+        })
+    }
+
     fn prepare_rename_impl(
         &mut self,
         buffer: Entity<Buffer>,

crates/proto/proto/zed.proto

@@ -327,7 +327,10 @@ message Envelope {
         Fetch fetch = 305;
         GetRemotes get_remotes = 306;
         GetRemotesResponse get_remotes_response = 307;
-        Pull pull = 308; // current max
+        Pull pull = 308;
+
+        ApplyCodeActionKind apply_code_action_kind = 309;
+        ApplyCodeActionKindResponse apply_code_action_kind_response = 310; // current max
     }
 
     reserved 87 to 88;
@@ -916,6 +919,16 @@ message ChannelBufferVersion {
     uint64 epoch = 3;
 }
 
+message ApplyCodeActionKind {
+    uint64 project_id = 1;
+    string kind = 2;
+    repeated uint64 buffer_ids = 3;
+}
+
+message ApplyCodeActionKindResponse {
+    ProjectTransaction transaction = 1;
+}
+
 enum FormatTrigger {
     Save = 0;
     Manual = 1;

crates/proto/src/proto.rs

@@ -236,6 +236,8 @@ messages!(
     (ExpandAllForProjectEntryResponse, Foreground),
     (Follow, Foreground),
     (FollowResponse, Foreground),
+    (ApplyCodeActionKind, Foreground),
+    (ApplyCodeActionKindResponse, Foreground),
     (FormatBuffers, Foreground),
     (FormatBuffersResponse, Foreground),
     (FuzzySearchUsers, Foreground),
@@ -472,6 +474,7 @@ request_messages!(
     (ExpandProjectEntry, ExpandProjectEntryResponse),
     (ExpandAllForProjectEntry, ExpandAllForProjectEntryResponse),
     (Follow, FollowResponse),
+    (ApplyCodeActionKind, ApplyCodeActionKindResponse),
     (FormatBuffers, FormatBuffersResponse),
     (FuzzySearchUsers, UsersResponse),
     (GetCachedEmbeddings, GetCachedEmbeddingsResponse),
@@ -610,6 +613,7 @@ entity_messages!(
     ExpandProjectEntry,
     ExpandAllForProjectEntry,
     FindSearchCandidates,
+    ApplyCodeActionKind,
     FormatBuffers,
     GetCodeActions,
     GetCompletions,