ruff server now supports commands for auto-fixing, organizing impor…

…ts, and formatting (astral-sh#10654)

## Summary

This builds off of the work in
astral-sh#10652 to implement a command
executor, backwards compatible with the commands from the previous LSP
(`ruff.applyAutofix`, `ruff.applyFormat` and
`ruff.applyOrganizeImports`).

This involved a lot of refactoring and tweaks to the code action
resolution code - the most notable change is that workspace edits are
specified in a slightly different way, using the more general `changes`
field instead of the `document_changes` field (which isn't supported on
all LSP clients). Additionally, the API for synchronous request handlers
has been updated to include access to the `Requester`, which we use to
send a `workspace/applyEdit` request to the client.

## Test Plan



https://github.com/astral-sh/ruff/assets/19577865/7932e30f-d944-4e35-b828-1d81aa56c087

crates/ruff_server/src/edit.rs

@@ -4,12 +4,16 @@ mod document;
 mod range;
 mod replacement;
 
+use std::collections::HashMap;
+
 pub use document::Document;
 pub(crate) use document::DocumentVersion;
 use lsp_types::PositionEncodingKind;
 pub(crate) use range::{RangeExt, ToRangeExt};
 pub(crate) use replacement::Replacement;
 
+use crate::session::ResolvedClientCapabilities;
+
 /// A convenient enumeration for supported text encodings. Can be converted to [`lsp_types::PositionEncodingKind`].
 // Please maintain the order from least to greatest priority for the derived `Ord` impl.
 #[derive(Default, Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
@@ -25,6 +29,14 @@ pub enum PositionEncoding {
     UTF8,
 }
 
+/// Tracks multi-document edits to eventually merge into a `WorkspaceEdit`.
+/// Compatible with clients that don't support `workspace.workspaceEdit.documentChanges`.
+#[derive(Debug)]
+pub(crate) enum WorkspaceEditTracker {
+    DocumentChanges(Vec<lsp_types::TextDocumentEdit>),
+    Changes(HashMap<lsp_types::Url, Vec<lsp_types::TextEdit>>),
+}
+
 impl From<PositionEncoding> for lsp_types::PositionEncodingKind {
     fn from(value: PositionEncoding) -> Self {
         match value {
@@ -50,3 +62,70 @@ impl TryFrom<&lsp_types::PositionEncodingKind> for PositionEncoding {
         })
     }
 }
+
+impl WorkspaceEditTracker {
+    pub(crate) fn new(client_capabilities: &ResolvedClientCapabilities) -> Self {
+        if client_capabilities.document_changes {
+            Self::DocumentChanges(Vec::default())
+        } else {
+            Self::Changes(HashMap::default())
+        }
+    }
+
+    /// Sets the edits made to a specific document. This should only be called
+    /// once for each document `uri`, and will fail if this is called for the same `uri`
+    /// multiple times.
+    pub(crate) fn set_edits_for_document(
+        &mut self,
+        uri: lsp_types::Url,
+        version: DocumentVersion,
+        edits: Vec<lsp_types::TextEdit>,
+    ) -> crate::Result<()> {
+        match self {
+            Self::DocumentChanges(document_edits) => {
+                if document_edits
+                    .iter()
+                    .any(|document| document.text_document.uri == uri)
+                {
+                    return Err(anyhow::anyhow!(
+                        "Attempted to add edits for a document that was already edited"
+                    ));
+                }
+                document_edits.push(lsp_types::TextDocumentEdit {
+                    text_document: lsp_types::OptionalVersionedTextDocumentIdentifier {
+                        uri,
+                        version: Some(version),
+                    },
+                    edits: edits.into_iter().map(lsp_types::OneOf::Left).collect(),
+                });
+                Ok(())
+            }
+            Self::Changes(changes) => {
+                if changes.get(&uri).is_some() {
+                    return Err(anyhow::anyhow!(
+                        "Attempted to add edits for a document that was already edited"
+                    ));
+                }
+                changes.insert(uri, edits);
+                Ok(())
+            }
+        }
+    }
+
+    pub(crate) fn is_empty(&self) -> bool {
+        match self {
+            Self::DocumentChanges(document_edits) => document_edits.is_empty(),
+            Self::Changes(changes) => changes.is_empty(),
+        }
+    }
+
+    pub(crate) fn into_workspace_edit(self) -> lsp_types::WorkspaceEdit {
+        match self {
+            Self::DocumentChanges(document_edits) => lsp_types::WorkspaceEdit {
+                document_changes: Some(lsp_types::DocumentChanges::Edits(document_edits)),
+                ..Default::default()
+            },
+            Self::Changes(changes) => lsp_types::WorkspaceEdit::new(changes),
+        }
+    }
+}

crates/ruff_server/src/lint.rs

@@ -35,7 +35,7 @@ pub(crate) struct DiagnosticFix {
     pub(crate) fixed_diagnostic: lsp_types::Diagnostic,
     pub(crate) title: String,
     pub(crate) code: String,
-    pub(crate) document_edits: Vec<lsp_types::TextDocumentEdit>,
+    pub(crate) edits: Vec<lsp_types::TextEdit>,
 }
 
 pub(crate) fn check(
@@ -90,11 +90,9 @@ pub(crate) fn check(
         .collect()
 }
 
-pub(crate) fn fixes_for_diagnostics<'d>(
-    document: &'d crate::edit::Document,
-    url: &'d lsp_types::Url,
+pub(crate) fn fixes_for_diagnostics(
+    document: &crate::edit::Document,
     encoding: PositionEncoding,
-    version: crate::edit::DocumentVersion,
     diagnostics: Vec<lsp_types::Diagnostic>,
 ) -> crate::Result<Vec<DiagnosticFix>> {
     diagnostics
@@ -118,22 +116,14 @@ pub(crate) fn fixes_for_diagnostics<'d>(
                         .to_range(document.contents(), document.index(), encoding),
                     new_text: edit.content().unwrap_or_default().to_string(),
                 });
-
-            let document_edits = vec![lsp_types::TextDocumentEdit {
-                text_document: lsp_types::OptionalVersionedTextDocumentIdentifier::new(
-                    url.clone(),
-                    version,
-                ),
-                edits: edits.map(lsp_types::OneOf::Left).collect(),
-            }];
             Ok(Some(DiagnosticFix {
                 fixed_diagnostic,
                 code: associated_data.code,
                 title: associated_data
                     .kind
                     .suggestion
                     .unwrap_or(associated_data.kind.name),
-                document_edits,
+                edits: edits.collect(),
             }))
         })
         .filter_map(crate::Result::transpose)

crates/ruff_server/src/server/api.rs

@@ -41,6 +41,7 @@ pub(super) fn request<'a>(req: server::Request) -> Task<'a> {
                 BackgroundSchedule::LatencySensitive,
             )
         }
+        request::ExecuteCommand::METHOD => local_request_task::<request::ExecuteCommand>(req),
         request::Format::METHOD => {
             background_request_task::<request::Format>(req, BackgroundSchedule::Fmt)
         }
@@ -87,13 +88,12 @@ pub(super) fn notification<'a>(notif: server::Notification) -> Task<'a> {
     })
 }
 
-#[allow(dead_code)]
 fn local_request_task<'a, R: traits::SyncRequestHandler>(
     req: server::Request,
 ) -> super::Result<Task<'a>> {
     let (id, params) = cast_request::<R>(req)?;
-    Ok(Task::local(|session, notifier, responder| {
-        let result = R::run(session, notifier, params);
+    Ok(Task::local(|session, notifier, requester, responder| {
+        let result = R::run(session, notifier, requester, params);
         respond::<R>(id, result, &responder);
     }))
 }
@@ -119,7 +119,7 @@ fn local_notification_task<'a, N: traits::SyncNotificationHandler>(
     notif: server::Notification,
 ) -> super::Result<Task<'a>> {
     let (id, params) = cast_notification::<N>(notif)?;
-    Ok(Task::local(move |session, notifier, _| {
+    Ok(Task::local(move |session, notifier, _, _| {
         if let Err(err) = N::run(session, notifier, params) {
             tracing::error!("An error occurred while running {id}: {err}");
         }

crates/ruff_server/src/server/api/requests.rs

@@ -1,16 +1,18 @@
 mod code_action;
 mod code_action_resolve;
 mod diagnostic;
+mod execute_command;
 mod format;
 mod format_range;
 
 use super::{
     define_document_url,
-    traits::{BackgroundDocumentRequestHandler, RequestHandler},
+    traits::{BackgroundDocumentRequestHandler, RequestHandler, SyncRequestHandler},
 };
 pub(super) use code_action::CodeActions;
 pub(super) use code_action_resolve::CodeActionResolve;
 pub(super) use diagnostic::DocumentDiagnostic;
+pub(super) use execute_command::ExecuteCommand;
 pub(super) use format::Format;
 pub(super) use format_range::FormatRange;
 

crates/ruff_server/src/server/api/requests/code_action.rs

@@ -1,3 +1,4 @@
+use crate::edit::WorkspaceEditTracker;
 use crate::lint::fixes_for_diagnostics;
 use crate::server::api::LSPResult;
 use crate::server::SupportedCodeAction;
@@ -50,30 +51,34 @@ impl super::BackgroundDocumentRequestHandler for CodeActions {
 fn quick_fix(
     snapshot: &DocumentSnapshot,
     diagnostics: Vec<types::Diagnostic>,
-) -> crate::Result<impl Iterator<Item = CodeActionOrCommand> + '_> {
+) -> crate::Result<Vec<CodeActionOrCommand>> {
     let document = snapshot.document();
 
-    let fixes = fixes_for_diagnostics(
-        document,
-        snapshot.url(),
-        snapshot.encoding(),
-        document.version(),
-        diagnostics,
-    )?;
-
-    Ok(fixes.into_iter().map(|fix| {
-        types::CodeActionOrCommand::CodeAction(types::CodeAction {
-            title: format!("{DIAGNOSTIC_NAME} ({}): {}", fix.code, fix.title),
-            kind: Some(types::CodeActionKind::QUICKFIX),
-            edit: Some(types::WorkspaceEdit {
-                document_changes: Some(types::DocumentChanges::Edits(fix.document_edits.clone())),
+    let fixes = fixes_for_diagnostics(document, snapshot.encoding(), diagnostics)?;
+
+    fixes
+        .into_iter()
+        .map(|fix| {
+            let mut tracker = WorkspaceEditTracker::new(snapshot.resolved_client_capabilities());
+
+            tracker.set_edits_for_document(
+                snapshot.url().clone(),
+                document.version(),
+                fix.edits,
+            )?;
+
+            Ok(types::CodeActionOrCommand::CodeAction(types::CodeAction {
+                title: format!("{DIAGNOSTIC_NAME} ({}): {}", fix.code, fix.title),
+                kind: Some(types::CodeActionKind::QUICKFIX),
+                edit: Some(tracker.into_workspace_edit()),
+                diagnostics: Some(vec![fix.fixed_diagnostic.clone()]),
+                data: Some(
+                    serde_json::to_value(snapshot.url()).expect("document url to serialize"),
+                ),
                 ..Default::default()
-            }),
-            diagnostics: Some(vec![fix.fixed_diagnostic.clone()]),
-            data: Some(serde_json::to_value(snapshot.url()).expect("document url to serialize")),
-            ..Default::default()
+            }))
         })
-    }))
+        .collect()
 }
 
 fn fix_all(snapshot: &DocumentSnapshot) -> crate::Result<CodeActionOrCommand> {
@@ -92,9 +97,11 @@ fn fix_all(snapshot: &DocumentSnapshot) -> crate::Result<CodeActionOrCommand> {
         (
             Some(resolve_edit_for_fix_all(
                 document,
+                snapshot.resolved_client_capabilities(),
                 snapshot.url(),
                 &snapshot.configuration().linter,
                 snapshot.encoding(),
+                document.version(),
             )?),
             None,
         )
@@ -125,9 +132,11 @@ fn organize_imports(snapshot: &DocumentSnapshot) -> crate::Result<CodeActionOrCo
         (
             Some(resolve_edit_for_organize_imports(
                 document,
+                snapshot.resolved_client_capabilities(),
                 snapshot.url(),
                 &snapshot.configuration().linter,
                 snapshot.encoding(),
+                document.version(),
             )?),
             None,
         )