Introduce notebook source actions

Author: snowsignal

crates/ruff_server/src/lib.rs

@@ -19,6 +19,10 @@ pub(crate) const DIAGNOSTIC_NAME: &str = "Ruff";
 pub(crate) const SOURCE_FIX_ALL_RUFF: CodeActionKind = CodeActionKind::new("source.fixAll.ruff");
 pub(crate) const SOURCE_ORGANIZE_IMPORTS_RUFF: CodeActionKind =
     CodeActionKind::new("source.organizeImports.ruff");
+pub(crate) const NOTEBOOK_SOURCE_FIX_ALL_RUFF: CodeActionKind =
+    CodeActionKind::new("notebook.source.fixAll.ruff");
+pub(crate) const NOTEBOOK_SOURCE_ORGANIZE_IMPORTS_RUFF: CodeActionKind =
+    CodeActionKind::new("notebook.source.organizeImports.ruff");
 
 /// A common result type used in most cases where a
 /// result type is needed.

crates/ruff_server/src/server.rs

@@ -292,8 +292,15 @@ pub(crate) enum SupportedCodeAction {
     SourceFixAll,
     /// Maps to `source.organizeImports` and `source.organizeImports.ruff` code action kinds.
     /// This is a source action that applies import sorting fixes to the currently open document.
-    #[allow(dead_code)] // TODO: remove
     SourceOrganizeImports,
+    /// Maps to the `notebook.source.fixAll` and `notebook.source.fixAll.ruff` code action kinds.
+    /// This is a source action, specifically for notebooks, that applies all safe fixes
+    /// to the currently open document.
+    NotebookSourceFixAll,
+    /// Maps to `source.organizeImports` and `source.organizeImports.ruff` code action kinds.
+    /// This is a source action, specifically for notebooks, that applies import sorting fixes
+    /// to the currently open document.
+    NotebookSourceOrganizeImports,
 }
 
 impl SupportedCodeAction {
@@ -303,6 +310,8 @@ impl SupportedCodeAction {
             Self::QuickFix => CodeActionKind::QUICKFIX,
             Self::SourceFixAll => crate::SOURCE_FIX_ALL_RUFF,
             Self::SourceOrganizeImports => crate::SOURCE_ORGANIZE_IMPORTS_RUFF,
+            Self::NotebookSourceFixAll => crate::NOTEBOOK_SOURCE_FIX_ALL_RUFF,
+            Self::NotebookSourceOrganizeImports => crate::NOTEBOOK_SOURCE_ORGANIZE_IMPORTS_RUFF,
         }
     }
 
@@ -318,6 +327,8 @@ impl SupportedCodeAction {
             Self::QuickFix,
             Self::SourceFixAll,
             Self::SourceOrganizeImports,
+            Self::NotebookSourceFixAll,
+            Self::NotebookSourceOrganizeImports,
         ]
         .into_iter()
     }

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

@@ -57,6 +57,20 @@ impl super::BackgroundDocumentRequestHandler for CodeActions {
             response.push(organize_imports(&snapshot).with_failure_code(ErrorCode::InternalError)?);
         }
 
+        if snapshot.client_settings().fix_all()
+            && supported_code_actions.contains(&SupportedCodeAction::NotebookSourceFixAll)
+        {
+            response.push(notebook_fix_all(&snapshot).with_failure_code(ErrorCode::InternalError)?);
+        }
+
+        if snapshot.client_settings().organize_imports()
+            && supported_code_actions.contains(&SupportedCodeAction::NotebookSourceOrganizeImports)
+        {
+            response.push(
+                notebook_organize_imports(&snapshot).with_failure_code(ErrorCode::InternalError)?,
+            );
+        }
+
         Ok(Some(response))
     }
 }
@@ -161,6 +175,42 @@ fn fix_all(snapshot: &DocumentSnapshot) -> crate::Result<CodeActionOrCommand> {
     }))
 }
 
+fn notebook_fix_all(snapshot: &DocumentSnapshot) -> crate::Result<CodeActionOrCommand> {
+    let document = snapshot.query();
+
+    let (edit, data) = if snapshot
+        .resolved_client_capabilities()
+        .code_action_deferred_edit_resolution
+    {
+        // The editor will request the edit in a `CodeActionsResolve` request
+        (
+            None,
+            Some(
+                serde_json::to_value(snapshot.query().make_key().into_url())
+                    .expect("document url should serialize"),
+            ),
+        )
+    } else {
+        (
+            Some(resolve_edit_for_fix_all(
+                document,
+                snapshot.resolved_client_capabilities(),
+                snapshot.query().settings().linter(),
+                snapshot.encoding(),
+            )?),
+            None,
+        )
+    };
+
+    Ok(CodeActionOrCommand::CodeAction(types::CodeAction {
+        title: format!("{DIAGNOSTIC_NAME}: Fix all auto-fixable problems"),
+        kind: Some(crate::NOTEBOOK_SOURCE_FIX_ALL_RUFF),
+        edit,
+        data,
+        ..Default::default()
+    }))
+}
+
 fn organize_imports(snapshot: &DocumentSnapshot) -> crate::Result<CodeActionOrCommand> {
     let document = snapshot.query();
 
@@ -197,6 +247,42 @@ fn organize_imports(snapshot: &DocumentSnapshot) -> crate::Result<CodeActionOrCo
     }))
 }
 
+fn notebook_organize_imports(snapshot: &DocumentSnapshot) -> crate::Result<CodeActionOrCommand> {
+    let document = snapshot.query();
+
+    let (edit, data) = if snapshot
+        .resolved_client_capabilities()
+        .code_action_deferred_edit_resolution
+    {
+        // The edit will be resolved later in the `CodeActionsResolve` request
+        (
+            None,
+            Some(
+                serde_json::to_value(snapshot.query().make_key().into_url())
+                    .expect("document url should serialize"),
+            ),
+        )
+    } else {
+        (
+            Some(resolve_edit_for_organize_imports(
+                document,
+                snapshot.resolved_client_capabilities(),
+                snapshot.query().settings().linter(),
+                snapshot.encoding(),
+            )?),
+            None,
+        )
+    };
+
+    Ok(CodeActionOrCommand::CodeAction(types::CodeAction {
+        title: format!("{DIAGNOSTIC_NAME}: Organize imports"),
+        kind: Some(crate::NOTEBOOK_SOURCE_ORGANIZE_IMPORTS_RUFF),
+        edit,
+        data,
+        ..Default::default()
+    }))
+}
+
 /// If `action_filter` is `None`, this returns [`SupportedCodeActionKind::all()`]. Otherwise,
 /// the list is filtered.
 fn supported_code_actions(

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

@@ -50,7 +50,7 @@ impl super::BackgroundDocumentRequestHandler for CodeActionResolve {
         };
 
         action.edit = match action_kind {
-            SupportedCodeAction::SourceFixAll => Some(
+            SupportedCodeAction::SourceFixAll | SupportedCodeAction::NotebookSourceFixAll => Some(
                 resolve_edit_for_fix_all(
                     query,
                     snapshot.resolved_client_capabilities(),
@@ -59,7 +59,8 @@ impl super::BackgroundDocumentRequestHandler for CodeActionResolve {
                 )
                 .with_failure_code(ErrorCode::InternalError)?,
             ),
-            SupportedCodeAction::SourceOrganizeImports => Some(
+            SupportedCodeAction::SourceOrganizeImports
+            | SupportedCodeAction::NotebookSourceOrganizeImports => Some(
                 resolve_edit_for_organize_imports(
                     query,
                     snapshot.resolved_client_capabilities(),