ruff server: Support Jupyter Notebook (*.ipynb) files (#11206)

## Summary

Closes https://github.com/astral-sh/ruff/issues/10858.

`ruff server` now supports `*.ipynb` (aka Jupyter Notebook) files.
Extensive internal changes have been made to facilitate this, which I've
done some work to contextualize with documentation and an pre-review
that highlights notable sections of the code.

`*.ipynb` cells should behave similarly to `*.py` documents, with one
major exception. The format command `ruff.applyFormat` will only apply
to the currently selected notebook cell - if you want to format an
entire notebook document, use `Format Notebook` from the VS Code context
menu.

## Test Plan

The VS Code extension does not yet have Jupyter Notebook support
enabled, so you'll first need to enable it manually. To do this,
checkout the `pre-release` branch and modify `src/common/server.ts` as
follows:

Before:
![Screenshot 2024-05-13 at 10 59
06 PM](c6a3c604-c405-4968-b8a2-5d670de89172)

After:
![Screenshot 2024-05-13 at 10 58
24 PM](94ab2e3d-0609-448d-9c8c-cd07c69a513b)

I recommend testing this PR with large, complicated notebook files. I
used notebook files from [this popular
repository](https://github.com/jakevdp/PythonDataScienceHandbook/tree/master/notebooks)
in my preliminary testing.

The main thing to test is ensuring that notebook cells behave the same
as Python documents, besides the aforementioned issue with
`ruff.applyFormat`. You should also test adding and deleting cells (in
particular, deleting all the code cells and ensure that doesn't break
anything), changing the kind of a cell (i.e. from markup -> code or vice
versa), and creating a new notebook file from scratch. Finally, you
should also test that source actions work as expected (and across the
entire notebook).

Note: `ruff.applyAutofix` and `ruff.applyOrganizeImports` are currently
broken for notebook files, and I suspect it has something to do with
https://github.com/astral-sh/ruff/issues/11248. Once this is fixed, I
will update the test plan accordingly.

---------

Co-authored-by: nolan <nolan.king90@gmail.com>

crates/ruff_server/src/server/api.rs

@@ -84,6 +84,15 @@ pub(super) fn notification<'a>(notif: server::Notification) -> Task<'a> {
         }
         notification::DidClose::METHOD => local_notification_task::<notification::DidClose>(notif),
         notification::DidOpen::METHOD => local_notification_task::<notification::DidOpen>(notif),
+        notification::DidOpenNotebook::METHOD => {
+            local_notification_task::<notification::DidOpenNotebook>(notif)
+        }
+        notification::DidChangeNotebook::METHOD => {
+            local_notification_task::<notification::DidChangeNotebook>(notif)
+        }
+        notification::DidCloseNotebook::METHOD => {
+            local_notification_task::<notification::DidCloseNotebook>(notif)
+        }
         method => {
             tracing::warn!("Received notification {method} which does not have a handler.");
             return Task::nothing();

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

@@ -1,17 +1,20 @@
-use crate::{server::client::Notifier, session::DocumentSnapshot};
+use crate::{
+    lint::Diagnostics,
+    server::client::Notifier,
+    session::{DocumentQuery, DocumentSnapshot},
+};
 
 use super::LSPResult;
 
-pub(super) fn generate_diagnostics(snapshot: &DocumentSnapshot) -> Vec<lsp_types::Diagnostic> {
+pub(super) fn generate_diagnostics(snapshot: &DocumentSnapshot) -> Diagnostics {
     if snapshot.client_settings().lint() {
         crate::lint::check(
-            snapshot.document(),
-            snapshot.url(),
-            snapshot.settings().linter(),
+            snapshot.query(),
+            snapshot.query().settings().linter(),
             snapshot.encoding(),
         )
     } else {
-        vec![]
+        Diagnostics::default()
     }
 }
 
@@ -19,31 +22,31 @@ pub(super) fn publish_diagnostics_for_document(
     snapshot: &DocumentSnapshot,
     notifier: &Notifier,
 ) -> crate::server::Result<()> {
-    let diagnostics = generate_diagnostics(snapshot);
-
-    notifier
-        .notify::<lsp_types::notification::PublishDiagnostics>(
-            lsp_types::PublishDiagnosticsParams {
-                uri: snapshot.url().clone(),
-                diagnostics,
-                version: Some(snapshot.document().version()),
-            },
-        )
-        .with_failure_code(lsp_server::ErrorCode::InternalError)?;
+    for (uri, diagnostics) in generate_diagnostics(snapshot) {
+        notifier
+            .notify::<lsp_types::notification::PublishDiagnostics>(
+                lsp_types::PublishDiagnosticsParams {
+                    uri,
+                    diagnostics,
+                    version: Some(snapshot.query().version()),
+                },
+            )
+            .with_failure_code(lsp_server::ErrorCode::InternalError)?;
+    }
 
     Ok(())
 }
 
 pub(super) fn clear_diagnostics_for_document(
-    snapshot: &DocumentSnapshot,
+    query: &DocumentQuery,
     notifier: &Notifier,
 ) -> crate::server::Result<()> {
     notifier
         .notify::<lsp_types::notification::PublishDiagnostics>(
             lsp_types::PublishDiagnosticsParams {
-                uri: snapshot.url().clone(),
+                uri: query.make_key().into_url(),
                 diagnostics: vec![],
-                version: Some(snapshot.document().version()),
+                version: Some(query.version()),
             },
         )
         .with_failure_code(lsp_server::ErrorCode::InternalError)?;

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

@@ -1,16 +1,22 @@
 mod cancel;
 mod did_change;
 mod did_change_configuration;
+mod did_change_notebook;
 mod did_change_watched_files;
 mod did_change_workspace;
 mod did_close;
+mod did_close_notebook;
 mod did_open;
+mod did_open_notebook;
 
 use super::traits::{NotificationHandler, SyncNotificationHandler};
 pub(super) use cancel::Cancel;
 pub(super) use did_change::DidChange;
 pub(super) use did_change_configuration::DidChangeConfiguration;
+pub(super) use did_change_notebook::DidChangeNotebook;
 pub(super) use did_change_watched_files::DidChangeWatchedFiles;
 pub(super) use did_change_workspace::DidChangeWorkspace;
 pub(super) use did_close::DidClose;
+pub(super) use did_close_notebook::DidCloseNotebook;
 pub(super) use did_open::DidOpen;
+pub(super) use did_open_notebook::DidOpenNotebook;

crates/ruff_server/src/server/api/notifications/cancel.rs

@@ -11,7 +11,6 @@ impl super::NotificationHandler for Cancel {
 }
 
 impl super::SyncNotificationHandler for Cancel {
-    #[tracing::instrument(skip_all)]
     fn run(
         _session: &mut Session,
         _notifier: Notifier,

crates/ruff_server/src/server/api/notifications/did_change.rs

@@ -3,6 +3,7 @@ use crate::server::api::LSPResult;
 use crate::server::client::{Notifier, Requester};
 use crate::server::Result;
 use crate::session::Session;
+use lsp_server::ErrorCode;
 use lsp_types as types;
 use lsp_types::notification as notif;
 
@@ -13,7 +14,6 @@ impl super::NotificationHandler for DidChange {
 }
 
 impl super::SyncNotificationHandler for DidChange {
-    #[tracing::instrument(skip_all, fields(file=%uri))]
     fn run(
         session: &mut Session,
         notifier: Notifier,
@@ -27,19 +27,13 @@ impl super::SyncNotificationHandler for DidChange {
             content_changes,
         }: types::DidChangeTextDocumentParams,
     ) -> Result<()> {
-        let encoding = session.encoding();
-        let document = session
-            .document_controller(&uri)
-            .with_failure_code(lsp_server::ErrorCode::InvalidParams)?;
+        let key = session
+            .key_from_url(&uri)
+            .with_failure_code(ErrorCode::InternalError)?;
 
-        if content_changes.is_empty() {
-            document.make_mut().update_version(new_version);
-            return Ok(());
-        }
-
-        document
-            .make_mut()
-            .apply_changes(content_changes, new_version, encoding);
+        session
+            .update_text_document(&key, content_changes, new_version)
+            .with_failure_code(ErrorCode::InternalError)?;
 
         // Publish diagnostics if the client doesnt support pull diagnostics
         if !session.resolved_client_capabilities().pull_diagnostics {

crates/ruff_server/src/server/api/notifications/did_change_notebook.rs

@@ -0,0 +1,41 @@
+use crate::server::api::diagnostics::publish_diagnostics_for_document;
+use crate::server::api::LSPResult;
+use crate::server::client::{Notifier, Requester};
+use crate::server::Result;
+use crate::session::Session;
+use lsp_server::ErrorCode;
+use lsp_types as types;
+use lsp_types::notification as notif;
+
+pub(crate) struct DidChangeNotebook;
+
+impl super::NotificationHandler for DidChangeNotebook {
+    type NotificationType = notif::DidChangeNotebookDocument;
+}
+
+impl super::SyncNotificationHandler for DidChangeNotebook {
+    fn run(
+        session: &mut Session,
+        notifier: Notifier,
+        _requester: &mut Requester,
+        types::DidChangeNotebookDocumentParams {
+            notebook_document: types::VersionedNotebookDocumentIdentifier { uri, version },
+            change: types::NotebookDocumentChangeEvent { cells, metadata },
+        }: types::DidChangeNotebookDocumentParams,
+    ) -> Result<()> {
+        let key = session
+            .key_from_url(&uri)
+            .with_failure_code(ErrorCode::InternalError)?;
+        session
+            .update_notebook_document(&key, cells, metadata, version)
+            .with_failure_code(ErrorCode::InternalError)?;
+
+        // publish new diagnostics
+        let snapshot = session
+            .take_snapshot(&uri)
+            .expect("snapshot should be available");
+        publish_diagnostics_for_document(&snapshot, &notifier)?;
+
+        Ok(())
+    }
+}

crates/ruff_server/src/server/api/notifications/did_change_watched_files.rs

@@ -20,9 +20,7 @@ impl super::SyncNotificationHandler for DidChangeWatchedFiles {
         params: types::DidChangeWatchedFilesParams,
     ) -> Result<()> {
         for change in &params.changes {
-            session
-                .reload_settings(&change.uri)
-                .with_failure_code(lsp_server::ErrorCode::InternalError)?;
+            session.reload_settings(&change.uri.to_file_path().unwrap());
         }
 
         if session.resolved_client_capabilities().workspace_refresh && !params.changes.is_empty() {

crates/ruff_server/src/server/api/notifications/did_change_workspace.rs

@@ -2,6 +2,8 @@ use crate::server::api::LSPResult;
 use crate::server::client::{Notifier, Requester};
 use crate::server::Result;
 use crate::session::Session;
+use anyhow::anyhow;
+use lsp_server::ErrorCode;
 use lsp_types as types;
 use lsp_types::notification as notif;
 
@@ -18,14 +20,21 @@ impl super::SyncNotificationHandler for DidChangeWorkspace {
         _requester: &mut Requester,
         params: types::DidChangeWorkspaceFoldersParams,
     ) -> Result<()> {
-        for new in params.event.added {
-            session
-                .open_workspace_folder(&new.uri)
-                .with_failure_code(lsp_server::ErrorCode::InvalidParams)?;
+        for types::WorkspaceFolder { ref uri, .. } in params.event.added {
+            let workspace_path = uri
+                .to_file_path()
+                .map_err(|()| anyhow!("expected document URI {uri} to be a valid file path"))
+                .with_failure_code(ErrorCode::InvalidParams)?;
+
+            session.open_workspace_folder(workspace_path);
         }
-        for removed in params.event.removed {
+        for types::WorkspaceFolder { ref uri, .. } in params.event.removed {
+            let workspace_path = uri
+                .to_file_path()
+                .map_err(|()| anyhow!("expected document URI {uri} to be a valid file path"))
+                .with_failure_code(ErrorCode::InvalidParams)?;
             session
-                .close_workspace_folder(&removed.uri)
+                .close_workspace_folder(&workspace_path)
                 .with_failure_code(lsp_server::ErrorCode::InvalidParams)?;
         }
         Ok(())

crates/ruff_server/src/server/api/notifications/did_close.rs

@@ -1,3 +1,4 @@
+use crate::edit::DocumentKey;
 use crate::server::api::diagnostics::clear_diagnostics_for_document;
 use crate::server::api::LSPResult;
 use crate::server::client::{Notifier, Requester};
@@ -13,7 +14,6 @@ impl super::NotificationHandler for DidClose {
 }
 
 impl super::SyncNotificationHandler for DidClose {
-    #[tracing::instrument(skip_all, fields(file=%uri))]
     fn run(
         session: &mut Session,
         notifier: Notifier,
@@ -22,20 +22,24 @@ impl super::SyncNotificationHandler for DidClose {
             text_document: types::TextDocumentIdentifier { uri },
         }: types::DidCloseTextDocumentParams,
     ) -> Result<()> {
-        // Publish an empty diagnostic report for the document if the client does not support pull diagnostics.
-        // This will de-register any existing diagnostics.
-        if !session.resolved_client_capabilities().pull_diagnostics {
-            let snapshot = session
-                .take_snapshot(&uri)
-                .ok_or_else(|| {
-                    anyhow::anyhow!("Unable to take snapshot for document with URL {uri}")
-                })
-                .with_failure_code(lsp_server::ErrorCode::InternalError)?;
-            clear_diagnostics_for_document(&snapshot, &notifier)?;
+        // Publish an empty diagnostic report for the document. This will de-register any existing diagnostics.
+        let snapshot = session
+            .take_snapshot(&uri)
+            .ok_or_else(|| anyhow::anyhow!("Unable to take snapshot for document with URL {uri}"))
+            .with_failure_code(lsp_server::ErrorCode::InternalError)?;
+        clear_diagnostics_for_document(snapshot.query(), &notifier)?;
+
+        let key = snapshot.query().make_key();
+
+        // Notebook cells will go through the `textDocument/didClose` path.
+        // We still want to publish empty diagnostics for them, but we
+        // shouldn't call `session.close_document` on them.
+        if matches!(key, DocumentKey::NotebookCell(_)) {
+            return Ok(());
         }
 
         session
-            .close_document(&uri)
+            .close_document(&key)
             .with_failure_code(lsp_server::ErrorCode::InternalError)
     }
 }

crates/ruff_server/src/server/api/notifications/did_close_notebook.rs

@@ -0,0 +1,35 @@
+use crate::server::api::LSPResult;
+use crate::server::client::{Notifier, Requester};
+use crate::server::Result;
+use crate::session::Session;
+use lsp_server::ErrorCode;
+use lsp_types as types;
+use lsp_types::notification as notif;
+
+pub(crate) struct DidCloseNotebook;
+
+impl super::NotificationHandler for DidCloseNotebook {
+    type NotificationType = notif::DidCloseNotebookDocument;
+}
+
+impl super::SyncNotificationHandler for DidCloseNotebook {
+    fn run(
+        session: &mut Session,
+        _notifier: Notifier,
+        _requester: &mut Requester,
+        types::DidCloseNotebookDocumentParams {
+            notebook_document: types::NotebookDocumentIdentifier { uri },
+            ..
+        }: types::DidCloseNotebookDocumentParams,
+    ) -> Result<()> {
+        let key = session
+            .key_from_url(&uri)
+            .with_failure_code(ErrorCode::InternalError)?;
+
+        session
+            .close_document(&key)
+            .with_failure_code(ErrorCode::InternalError)?;
+
+        Ok(())
+    }
+}

crates/ruff_server/src/server/api/notifications/did_open.rs

@@ -3,6 +3,9 @@ use crate::server::api::LSPResult;
 use crate::server::client::{Notifier, Requester};
 use crate::server::Result;
 use crate::session::Session;
+use crate::TextDocument;
+use anyhow::anyhow;
+use lsp_server::ErrorCode;
 use lsp_types as types;
 use lsp_types::notification as notif;
 
@@ -13,7 +16,6 @@ impl super::NotificationHandler for DidOpen {
 }
 
 impl super::SyncNotificationHandler for DidOpen {
-    #[tracing::instrument(skip_all, fields(file=%url))]
     fn run(
         session: &mut Session,
         notifier: Notifier,
@@ -21,21 +23,28 @@ impl super::SyncNotificationHandler for DidOpen {
         types::DidOpenTextDocumentParams {
             text_document:
                 types::TextDocumentItem {
-                    uri: ref url,
+                    ref uri,
                     text,
                     version,
                     ..
                 },
         }: types::DidOpenTextDocumentParams,
     ) -> Result<()> {
-        session.open_document(url, text, version);
+        let document_path: std::path::PathBuf = uri
+            .to_file_path()
+            .map_err(|()| anyhow!("expected document URI {uri} to be a valid file path"))
+            .with_failure_code(ErrorCode::InvalidParams)?;
+
+        let document = TextDocument::new(text, version);
+
+        session.open_text_document(document_path, document);
 
         // Publish diagnostics if the client doesnt support pull diagnostics
         if !session.resolved_client_capabilities().pull_diagnostics {
             let snapshot = session
-                .take_snapshot(url)
+                .take_snapshot(uri)
                 .ok_or_else(|| {
-                    anyhow::anyhow!("Unable to take snapshot for document with URL {url}")
+                    anyhow::anyhow!("Unable to take snapshot for document with URL {uri}")
                 })
                 .with_failure_code(lsp_server::ErrorCode::InternalError)?;
             publish_diagnostics_for_document(&snapshot, &notifier)?;

crates/ruff_server/src/server/api/notifications/did_open_notebook.rs

@@ -0,0 +1,58 @@
+use crate::edit::NotebookDocument;
+use crate::server::api::diagnostics::publish_diagnostics_for_document;
+use crate::server::api::LSPResult;
+use crate::server::client::{Notifier, Requester};
+use crate::server::Result;
+use crate::session::Session;
+use anyhow::anyhow;
+use lsp_server::ErrorCode;
+use lsp_types as types;
+use lsp_types::notification as notif;
+
+pub(crate) struct DidOpenNotebook;
+
+impl super::NotificationHandler for DidOpenNotebook {
+    type NotificationType = notif::DidOpenNotebookDocument;
+}
+
+impl super::SyncNotificationHandler for DidOpenNotebook {
+    fn run(
+        session: &mut Session,
+        notifier: Notifier,
+        _requester: &mut Requester,
+        types::DidOpenNotebookDocumentParams {
+            notebook_document:
+                types::NotebookDocument {
+                    uri,
+                    version,
+                    cells,
+                    metadata,
+                    ..
+                },
+            cell_text_documents,
+        }: types::DidOpenNotebookDocumentParams,
+    ) -> Result<()> {
+        let notebook = NotebookDocument::new(
+            version,
+            cells,
+            metadata.unwrap_or_default(),
+            cell_text_documents,
+        )
+        .with_failure_code(ErrorCode::InternalError)?;
+
+        let notebook_path = uri
+            .to_file_path()
+            .map_err(|()| anyhow!("expected notebook URI {uri} to be a valid file path"))
+            .with_failure_code(ErrorCode::InvalidParams)?;
+
+        session.open_notebook_document(notebook_path, notebook);
+
+        // publish diagnostics
+        let snapshot = session
+            .take_snapshot(&uri)
+            .expect("snapshot should be available");
+        publish_diagnostics_for_document(&snapshot, &notifier)?;
+
+        Ok(())
+    }
+}

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

@@ -29,12 +29,8 @@ impl super::BackgroundDocumentRequestHandler for CodeActions {
 
         let supported_code_actions = supported_code_actions(params.context.only.clone());
 
-        let fixes = fixes_for_diagnostics(
-            snapshot.document(),
-            snapshot.encoding(),
-            params.context.diagnostics,
-        )
-        .with_failure_code(ErrorCode::InternalError)?;
+        let fixes = fixes_for_diagnostics(params.context.diagnostics)
+            .with_failure_code(ErrorCode::InternalError)?;
 
         if snapshot.client_settings().fix_violation()
             && supported_code_actions.contains(&SupportedCodeAction::QuickFix)
@@ -61,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))
     }
 }
@@ -69,7 +79,8 @@ fn quick_fix(
     snapshot: &DocumentSnapshot,
     fixes: &[DiagnosticFix],
 ) -> crate::Result<Vec<CodeActionOrCommand>> {
-    let document = snapshot.document();
+    let document = snapshot.query();
+
     fixes
         .iter()
         .filter(|fix| !fix.edits.is_empty())
@@ -77,7 +88,7 @@ fn quick_fix(
             let mut tracker = WorkspaceEditTracker::new(snapshot.resolved_client_capabilities());
 
             tracker.set_edits_for_document(
-                snapshot.url().clone(),
+                snapshot.query().make_key().into_url(),
                 document.version(),
                 fix.edits.clone(),
             )?;
@@ -88,7 +99,8 @@ fn quick_fix(
                 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"),
+                    serde_json::to_value(snapshot.query().make_key().into_url())
+                        .expect("document url should serialize"),
                 ),
                 ..Default::default()
             }))
@@ -106,8 +118,8 @@ fn noqa_comments(snapshot: &DocumentSnapshot, fixes: &[DiagnosticFix]) -> Vec<Co
 
             tracker
                 .set_edits_for_document(
-                    snapshot.url().clone(),
-                    snapshot.document().version(),
+                    snapshot.query().make_key().into_url(),
+                    snapshot.query().version(),
                     vec![edit],
                 )
                 .ok()?;
@@ -118,7 +130,8 @@ fn noqa_comments(snapshot: &DocumentSnapshot, fixes: &[DiagnosticFix]) -> Vec<Co
                 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"),
+                    serde_json::to_value(snapshot.query().make_key().into_url())
+                        .expect("document url should serialize"),
                 ),
                 ..Default::default()
             }))
@@ -127,7 +140,7 @@ fn noqa_comments(snapshot: &DocumentSnapshot, fixes: &[DiagnosticFix]) -> Vec<Co
 }
 
 fn fix_all(snapshot: &DocumentSnapshot) -> crate::Result<CodeActionOrCommand> {
-    let document = snapshot.document();
+    let document = snapshot.query();
 
     let (edit, data) = if snapshot
         .resolved_client_capabilities()
@@ -136,17 +149,18 @@ fn fix_all(snapshot: &DocumentSnapshot) -> crate::Result<CodeActionOrCommand> {
         // The editor will request the edit in a `CodeActionsResolve` request
         (
             None,
-            Some(serde_json::to_value(snapshot.url()).expect("document url to serialize")),
+            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.url(),
-                snapshot.settings().linter(),
+                snapshot.query().settings().linter(),
                 snapshot.encoding(),
-                document.version(),
             )?),
             None,
         )
@@ -161,8 +175,44 @@ 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.document();
+    let document = snapshot.query();
 
     let (edit, data) = if snapshot
         .resolved_client_capabilities()
@@ -171,17 +221,18 @@ fn organize_imports(snapshot: &DocumentSnapshot) -> crate::Result<CodeActionOrCo
         // The edit will be resolved later in the `CodeActionsResolve` request
         (
             None,
-            Some(serde_json::to_value(snapshot.url()).expect("document url to serialize")),
+            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.url(),
-                snapshot.settings().linter(),
+                snapshot.query().settings().linter(),
                 snapshot.encoding(),
-                document.version(),
             )?),
             None,
         )
@@ -196,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

@@ -1,10 +1,11 @@
 use std::borrow::Cow;
 
-use crate::edit::{DocumentVersion, WorkspaceEditTracker};
+use crate::edit::WorkspaceEditTracker;
+use crate::fix::Fixes;
 use crate::server::api::LSPResult;
 use crate::server::SupportedCodeAction;
 use crate::server::{client::Notifier, Result};
-use crate::session::{DocumentSnapshot, ResolvedClientCapabilities};
+use crate::session::{DocumentQuery, DocumentSnapshot, ResolvedClientCapabilities};
 use crate::PositionEncoding;
 use lsp_server::ErrorCode;
 use lsp_types::{self as types, request as req};
@@ -28,7 +29,7 @@ impl super::BackgroundDocumentRequestHandler for CodeActionResolve {
         _notifier: Notifier,
         mut action: types::CodeAction,
     ) -> Result<types::CodeAction> {
-        let document = snapshot.document();
+        let query = snapshot.query();
 
         let code_actions = SupportedCodeAction::from_kind(
             action
@@ -49,25 +50,22 @@ impl super::BackgroundDocumentRequestHandler for CodeActionResolve {
         };
 
         action.edit = match action_kind {
-            SupportedCodeAction::SourceFixAll => Some(
+            SupportedCodeAction::SourceFixAll | SupportedCodeAction::NotebookSourceFixAll => Some(
                 resolve_edit_for_fix_all(
-                    document,
+                    query,
                     snapshot.resolved_client_capabilities(),
-                    snapshot.url(),
-                    snapshot.settings().linter(),
+                    query.settings().linter(),
                     snapshot.encoding(),
-                    document.version(),
                 )
                 .with_failure_code(ErrorCode::InternalError)?,
             ),
-            SupportedCodeAction::SourceOrganizeImports => Some(
+            SupportedCodeAction::SourceOrganizeImports
+            | SupportedCodeAction::NotebookSourceOrganizeImports => Some(
                 resolve_edit_for_organize_imports(
-                    document,
+                    query,
                     snapshot.resolved_client_capabilities(),
-                    snapshot.url(),
-                    snapshot.settings().linter(),
+                    query.settings().linter(),
                     snapshot.encoding(),
-                    document.version(),
                 )
                 .with_failure_code(ErrorCode::InternalError)?,
             ),
@@ -84,54 +82,46 @@ impl super::BackgroundDocumentRequestHandler for CodeActionResolve {
 }
 
 pub(super) fn resolve_edit_for_fix_all(
-    document: &crate::edit::Document,
+    query: &DocumentQuery,
     client_capabilities: &ResolvedClientCapabilities,
-    url: &types::Url,
     linter_settings: &LinterSettings,
     encoding: PositionEncoding,
-    version: DocumentVersion,
 ) -> crate::Result<types::WorkspaceEdit> {
     let mut tracker = WorkspaceEditTracker::new(client_capabilities);
-    tracker.set_edits_for_document(
-        url.clone(),
-        version,
-        fix_all_edit(document, url, linter_settings, encoding)?,
+    tracker.set_fixes_for_document(
+        fix_all_edit(query, linter_settings, encoding)?,
+        query.version(),
     )?;
     Ok(tracker.into_workspace_edit())
 }
 
 pub(super) fn fix_all_edit(
-    document: &crate::edit::Document,
-    document_url: &types::Url,
+    query: &DocumentQuery,
     linter_settings: &LinterSettings,
     encoding: PositionEncoding,
-) -> crate::Result<Vec<types::TextEdit>> {
-    crate::fix::fix_all(document, document_url, linter_settings, encoding)
+) -> crate::Result<Fixes> {
+    crate::fix::fix_all(query, linter_settings, encoding)
 }
 
 pub(super) fn resolve_edit_for_organize_imports(
-    document: &crate::edit::Document,
+    query: &DocumentQuery,
     client_capabilities: &ResolvedClientCapabilities,
-    url: &types::Url,
     linter_settings: &ruff_linter::settings::LinterSettings,
     encoding: PositionEncoding,
-    version: DocumentVersion,
 ) -> crate::Result<types::WorkspaceEdit> {
     let mut tracker = WorkspaceEditTracker::new(client_capabilities);
-    tracker.set_edits_for_document(
-        url.clone(),
-        version,
-        organize_imports_edit(document, url, linter_settings, encoding)?,
+    tracker.set_fixes_for_document(
+        organize_imports_edit(query, linter_settings, encoding)?,
+        query.version(),
     )?;
     Ok(tracker.into_workspace_edit())
 }
 
 pub(super) fn organize_imports_edit(
-    document: &crate::edit::Document,
-    document_url: &types::Url,
+    query: &DocumentQuery,
     linter_settings: &LinterSettings,
     encoding: PositionEncoding,
-) -> crate::Result<Vec<types::TextEdit>> {
+) -> crate::Result<Fixes> {
     let mut linter_settings = linter_settings.clone();
     linter_settings.rules = [
         Rule::UnsortedImports,       // I001
@@ -140,5 +130,5 @@ pub(super) fn organize_imports_edit(
     .into_iter()
     .collect();
 
-    crate::fix::fix_all(document, document_url, &linter_settings, encoding)
+    crate::fix::fix_all(query, &linter_settings, encoding)
 }

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

@@ -26,7 +26,13 @@ impl super::BackgroundDocumentRequestHandler for DocumentDiagnostic {
                 full_document_diagnostic_report: FullDocumentDiagnosticReport {
                     // TODO(jane): eventually this will be important for caching diagnostic information.
                     result_id: None,
-                    items: generate_diagnostics(&snapshot),
+                    // Pull diagnostic requests are only called for text documents.
+                    // Since diagnostic requests generate
+                    items: generate_diagnostics(&snapshot)
+                        .into_iter()
+                        .next()
+                        .map(|(_, diagnostics)| diagnostics)
+                        .unwrap_or_default(),
                 },
             }),
         ))

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

@@ -21,7 +21,7 @@ enum Command {
 pub(crate) struct ExecuteCommand;
 
 #[derive(Deserialize)]
-struct TextDocumentArgument {
+struct Argument {
     uri: types::Url,
     version: DocumentVersion,
 }
@@ -45,7 +45,7 @@ impl super::SyncRequestHandler for ExecuteCommand {
             return Err(anyhow::anyhow!("Cannot execute the '{}' command: the client does not support `workspace/applyEdit`", command.label())).with_failure_code(ErrorCode::InternalError);
         }
 
-        let mut arguments: Vec<TextDocumentArgument> = params
+        let mut arguments: Vec<Argument> = params
             .arguments
             .into_iter()
             .map(|value| serde_json::from_value(value).with_failure_code(ErrorCode::InvalidParams))
@@ -55,22 +55,21 @@ impl super::SyncRequestHandler for ExecuteCommand {
         arguments.dedup_by(|a, b| a.uri == b.uri);
 
         let mut edit_tracker = WorkspaceEditTracker::new(session.resolved_client_capabilities());
-        for TextDocumentArgument { uri, version } in arguments {
+        for Argument { uri, version } in arguments {
             let snapshot = session
                 .take_snapshot(&uri)
                 .ok_or(anyhow::anyhow!("Document snapshot not available for {uri}",))
                 .with_failure_code(ErrorCode::InternalError)?;
             match command {
                 Command::FixAll => {
-                    let edits = super::code_action_resolve::fix_all_edit(
-                        snapshot.document(),
-                        snapshot.url(),
-                        snapshot.settings().linter(),
+                    let fixes = super::code_action_resolve::fix_all_edit(
+                        snapshot.query(),
+                        snapshot.query().settings().linter(),
                         snapshot.encoding(),
                     )
                     .with_failure_code(ErrorCode::InternalError)?;
                     edit_tracker
-                        .set_edits_for_document(uri, version, edits)
+                        .set_fixes_for_document(fixes, snapshot.query().version())
                         .with_failure_code(ErrorCode::InternalError)?;
                 }
                 Command::Format => {
@@ -82,15 +81,14 @@ impl super::SyncRequestHandler for ExecuteCommand {
                     }
                 }
                 Command::OrganizeImports => {
-                    let edits = super::code_action_resolve::organize_imports_edit(
-                        snapshot.document(),
-                        snapshot.url(),
-                        snapshot.settings().linter(),
+                    let fixes = super::code_action_resolve::organize_imports_edit(
+                        snapshot.query(),
+                        snapshot.query().settings().linter(),
                         snapshot.encoding(),
                     )
                     .with_failure_code(ErrorCode::InternalError)?;
                     edit_tracker
-                        .set_edits_for_document(uri, version, edits)
+                        .set_fixes_for_document(fixes, snapshot.query().version())
                         .with_failure_code(ErrorCode::InternalError)?;
                 }
             }

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

@@ -24,14 +24,37 @@ impl super::BackgroundDocumentRequestHandler for Format {
 }
 
 pub(super) fn format_document(snapshot: &DocumentSnapshot) -> Result<super::FormatResponse> {
-    let doc = snapshot.document();
+    let doc = snapshot
+        .query()
+        .as_single_document()
+        .expect("format should only be called on text documents or notebook cells");
     let source = doc.contents();
-    let formatted = crate::format::format(doc, snapshot.settings().formatter())
-        .with_failure_code(lsp_server::ErrorCode::InternalError)?;
+    let mut formatted = crate::format::format(
+        doc,
+        snapshot.query().source_type(),
+        snapshot.query().settings().formatter(),
+    )
+    .with_failure_code(lsp_server::ErrorCode::InternalError)?;
     // fast path - if the code is the same, return early
     if formatted == source {
         return Ok(None);
     }
+
+    // special case - avoid adding a newline to a notebook cell if it didn't already exist
+    if snapshot.query().as_notebook().is_some() {
+        let mut trimmed = formatted.as_str();
+        if !source.ends_with("\r\n") {
+            trimmed = trimmed.trim_end_matches("\r\n");
+        }
+        if !source.ends_with('\n') {
+            trimmed = trimmed.trim_end_matches('\n');
+        }
+        if !source.ends_with('\r') {
+            trimmed = trimmed.trim_end_matches('\r');
+        }
+        formatted = trimmed.to_string();
+    }
+
     let formatted_index: LineIndex = LineIndex::from_source_text(&formatted);
 
     let unformatted_index = doc.index();

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

@@ -17,13 +17,20 @@ impl super::BackgroundDocumentRequestHandler for FormatRange {
         _notifier: Notifier,
         params: types::DocumentRangeFormattingParams,
     ) -> Result<super::FormatResponse> {
-        let document = snapshot.document();
+        let document = snapshot
+            .query()
+            .as_single_document()
+            .expect("hover should only be called on text documents or notebook cells");
         let text = document.contents();
         let index = document.index();
         let range = params.range.to_text_range(text, index, snapshot.encoding());
-        let formatted_range =
-            crate::format::format_range(document, snapshot.settings().formatter(), range)
-                .with_failure_code(lsp_server::ErrorCode::InternalError)?;
+        let formatted_range = crate::format::format_range(
+            document,
+            snapshot.query().source_type(),
+            snapshot.query().settings().formatter(),
+            range,
+        )
+        .with_failure_code(lsp_server::ErrorCode::InternalError)?;
         Ok(Some(vec![types::TextEdit {
             range: formatted_range
                 .source_range()

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

@@ -29,7 +29,11 @@ pub(crate) fn hover(
     snapshot: &DocumentSnapshot,
     position: &types::TextDocumentPositionParams,
 ) -> Option<types::Hover> {
-    let document = snapshot.document();
+    // Hover only operates on text documents or notebook cells
+    let document = snapshot
+        .query()
+        .as_single_document()
+        .expect("hover should only be called on text documents or notebook cells");
     let line_number: usize = position
         .position
         .line