[red-knot] Incorporate recent ruff server improvements into red knot's LSP (#17044)

crates/red_knot_server/src/edit/notebook.rs

@@ -136,17 +136,15 @@ impl NotebookDocument {
                 // provide the actual contents of the cells, so we'll initialize them with empty
                 // contents.
                 for cell in structure.array.cells.into_iter().flatten().rev() {
-                    if let Some(text_document) = deleted_cells.remove(&cell.document) {
-                        let version = text_document.version();
-                        self.cells.push(NotebookCell::new(
-                            cell,
-                            text_document.into_contents(),
-                            version,
-                        ));
-                    } else {
-                        self.cells
-                            .insert(start, NotebookCell::new(cell, String::new(), 0));
-                    }
+                    let (content, version) =
+                        if let Some(text_document) = deleted_cells.remove(&cell.document) {
+                            let version = text_document.version();
+                            (text_document.into_contents(), version)
+                        } else {
+                            (String::new(), 0)
+                        };
+                    self.cells
+                        .insert(start, NotebookCell::new(cell, content, version));
                 }
 
                 // Third, register the new cells in the index and update existing ones that came
@@ -200,6 +198,11 @@ impl NotebookDocument {
         self.version
     }
 
+    /// Get the URI for a cell by its index within the cell array.
+    pub(crate) fn cell_uri_by_index(&self, index: CellId) -> Option<&lsp_types::Url> {
+        self.cells.get(index).map(|cell| &cell.url)
+    }
+
     /// Get the text document representing the contents of a cell by the cell URI.
     pub(crate) fn cell_document_by_uri(&self, uri: &lsp_types::Url) -> Option<&TextDocument> {
         self.cells
@@ -238,3 +241,115 @@ impl NotebookCell {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::NotebookDocument;
+
+    enum TestCellContent {
+        #[allow(dead_code)]
+        Markup(String),
+        Code(String),
+    }
+
+    fn create_test_url(index: usize) -> lsp_types::Url {
+        lsp_types::Url::parse(&format!("cell:/test.ipynb#{index}")).unwrap()
+    }
+
+    fn create_test_notebook(test_cells: Vec<TestCellContent>) -> NotebookDocument {
+        let mut cells = Vec::with_capacity(test_cells.len());
+        let mut cell_documents = Vec::with_capacity(test_cells.len());
+
+        for (index, test_cell) in test_cells.into_iter().enumerate() {
+            let url = create_test_url(index);
+            match test_cell {
+                TestCellContent::Markup(content) => {
+                    cells.push(lsp_types::NotebookCell {
+                        kind: lsp_types::NotebookCellKind::Markup,
+                        document: url.clone(),
+                        metadata: None,
+                        execution_summary: None,
+                    });
+                    cell_documents.push(lsp_types::TextDocumentItem {
+                        uri: url,
+                        language_id: "markdown".to_owned(),
+                        version: 0,
+                        text: content,
+                    });
+                }
+                TestCellContent::Code(content) => {
+                    cells.push(lsp_types::NotebookCell {
+                        kind: lsp_types::NotebookCellKind::Code,
+                        document: url.clone(),
+                        metadata: None,
+                        execution_summary: None,
+                    });
+                    cell_documents.push(lsp_types::TextDocumentItem {
+                        uri: url,
+                        language_id: "python".to_owned(),
+                        version: 0,
+                        text: content,
+                    });
+                }
+            }
+        }
+
+        NotebookDocument::new(0, cells, serde_json::Map::default(), cell_documents).unwrap()
+    }
+
+    /// This test case checks that for a notebook with three code cells, when the client sends a
+    /// change request to swap the first two cells, the notebook document is updated correctly.
+    ///
+    /// The swap operation as a change request is represented as deleting the first two cells and
+    /// adding them back in the reverse order.
+    #[test]
+    fn swap_cells() {
+        let mut notebook = create_test_notebook(vec![
+            TestCellContent::Code("cell = 0".to_owned()),
+            TestCellContent::Code("cell = 1".to_owned()),
+            TestCellContent::Code("cell = 2".to_owned()),
+        ]);
+
+        notebook
+            .update(
+                Some(lsp_types::NotebookDocumentCellChange {
+                    structure: Some(lsp_types::NotebookDocumentCellChangeStructure {
+                        array: lsp_types::NotebookCellArrayChange {
+                            start: 0,
+                            delete_count: 2,
+                            cells: Some(vec![
+                                lsp_types::NotebookCell {
+                                    kind: lsp_types::NotebookCellKind::Code,
+                                    document: create_test_url(1),
+                                    metadata: None,
+                                    execution_summary: None,
+                                },
+                                lsp_types::NotebookCell {
+                                    kind: lsp_types::NotebookCellKind::Code,
+                                    document: create_test_url(0),
+                                    metadata: None,
+                                    execution_summary: None,
+                                },
+                            ]),
+                        },
+                        did_open: None,
+                        did_close: None,
+                    }),
+                    data: None,
+                    text_content: None,
+                }),
+                None,
+                1,
+                crate::PositionEncoding::default(),
+            )
+            .unwrap();
+
+        assert_eq!(
+            notebook.make_ruff_notebook().source_code(),
+            "cell = 1
+cell = 0
+cell = 2
+"
+        );
+    }
+}

crates/red_knot_server/src/edit/text_document.rs

@@ -20,6 +20,23 @@ pub struct TextDocument {
     /// The latest version of the document, set by the LSP client. The server will panic in
     /// debug mode if we attempt to update the document with an 'older' version.
     version: DocumentVersion,
+    /// The language ID of the document as provided by the client.
+    language_id: Option<LanguageId>,
+}
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum LanguageId {
+    Python,
+    Other,
+}
+
+impl From<&str> for LanguageId {
+    fn from(language_id: &str) -> Self {
+        match language_id {
+            "python" => Self::Python,
+            _ => Self::Other,
+        }
+    }
 }
 
 impl TextDocument {
@@ -29,9 +46,16 @@ impl TextDocument {
             contents,
             index,
             version,
+            language_id: None,
         }
     }
 
+    #[must_use]
+    pub fn with_language_id(mut self, language_id: &str) -> Self {
+        self.language_id = Some(LanguageId::from(language_id));
+        self
+    }
+
     pub fn into_contents(self) -> String {
         self.contents
     }
@@ -48,6 +72,10 @@ impl TextDocument {
         self.version
     }
 
+    pub fn language_id(&self) -> Option<LanguageId> {
+        self.language_id
+    }
+
     pub fn apply_changes(
         &mut self,
         changes: Vec<lsp_types::TextDocumentContentChangeEvent>,
@@ -66,7 +94,6 @@ impl TextDocument {
             return;
         }
 
-        let old_contents = self.contents().to_string();
         let mut new_contents = self.contents().to_string();
         let mut active_index = self.index().clone();
 
@@ -87,15 +114,11 @@ impl TextDocument {
                 new_contents = change;
             }
 
-            if new_contents != old_contents {
-                active_index = LineIndex::from_source_text(&new_contents);
-            }
+            active_index = LineIndex::from_source_text(&new_contents);
         }
 
         self.modify_with_manual_index(|contents, version, index| {
-            if contents != &new_contents {
-                *index = active_index;
-            }
+            *index = active_index;
             *contents = new_contents;
             *version = new_version;
         });
@@ -125,3 +148,75 @@ impl TextDocument {
         debug_assert!(self.version >= old_version);
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::{PositionEncoding, TextDocument};
+    use lsp_types::{Position, TextDocumentContentChangeEvent};
+
+    #[test]
+    fn redo_edit() {
+        let mut document = TextDocument::new(
+            r#""""
+测试comment
+一些测试内容
+"""
+import click
+
+
+@click.group()
+def interface():
+    pas
+"#
+            .to_string(),
+            0,
+        );
+
+        // Add an `s`, remove it again (back to the original code), and then re-add the `s`
+        document.apply_changes(
+            vec![
+                TextDocumentContentChangeEvent {
+                    range: Some(lsp_types::Range::new(
+                        Position::new(9, 7),
+                        Position::new(9, 7),
+                    )),
+                    range_length: Some(0),
+                    text: "s".to_string(),
+                },
+                TextDocumentContentChangeEvent {
+                    range: Some(lsp_types::Range::new(
+                        Position::new(9, 7),
+                        Position::new(9, 8),
+                    )),
+                    range_length: Some(1),
+                    text: String::new(),
+                },
+                TextDocumentContentChangeEvent {
+                    range: Some(lsp_types::Range::new(
+                        Position::new(9, 7),
+                        Position::new(9, 7),
+                    )),
+                    range_length: Some(0),
+                    text: "s".to_string(),
+                },
+            ],
+            1,
+            PositionEncoding::UTF16,
+        );
+
+        assert_eq!(
+            &document.contents,
+            r#""""
+测试comment
+一些测试内容
+"""
+import click
+
+
+@click.group()
+def interface():
+    pass
+"#
+        );
+    }
+}

crates/red_knot_server/src/server/api.rs

@@ -1,8 +1,8 @@
-use lsp_server as server;
-
 use crate::server::schedule::Task;
 use crate::session::Session;
 use crate::system::{url_to_any_system_path, AnySystemPath};
+use lsp_server as server;
+use lsp_types::notification::Notification;
 
 mod diagnostics;
 mod notifications;
@@ -52,6 +52,11 @@ pub(super) fn notification<'a>(notif: server::Notification) -> Task<'a> {
         notification::DidCloseNotebookHandler::METHOD => {
             local_notification_task::<notification::DidCloseNotebookHandler>(notif)
         }
+        lsp_types::notification::SetTrace::METHOD => {
+            tracing::trace!("Ignoring `setTrace` notification");
+            return Task::nothing();
+        },
+
         method => {
             tracing::warn!("Received notification {method} which does not have a handler.");
             return Task::nothing();
@@ -69,6 +74,7 @@ fn _local_request_task<'a, R: traits::SyncRequestHandler>(
 ) -> super::Result<Task<'a>> {
     let (id, params) = cast_request::<R>(req)?;
     Ok(Task::local(|session, notifier, requester, responder| {
+        let _span = tracing::trace_span!("request", %id, method = R::METHOD).entered();
         let result = R::run(session, notifier, requester, params);
         respond::<R>(id, result, &responder);
     }))
@@ -98,6 +104,7 @@ fn background_request_task<'a, R: traits::BackgroundDocumentRequestHandler>(
         };
 
         Box::new(move |notifier, responder| {
+            let _span = tracing::trace_span!("request", %id, method = R::METHOD).entered();
             let result = R::run_with_snapshot(snapshot, db, notifier, params);
             respond::<R>(id, result, &responder);
         })
@@ -109,6 +116,7 @@ fn local_notification_task<'a, N: traits::SyncNotificationHandler>(
 ) -> super::Result<Task<'a>> {
     let (id, params) = cast_notification::<N>(notif)?;
     Ok(Task::local(move |session, notifier, requester, _| {
+        let _span = tracing::trace_span!("notification", method = N::METHOD).entered();
         if let Err(err) = N::run(session, notifier, requester, params) {
             tracing::error!("An error occurred while running {id}: {err}");
             show_err_msg!("Ruff encountered a problem. Check the logs for more details.");
@@ -128,6 +136,7 @@ fn background_notification_thread<'a, N: traits::BackgroundDocumentNotificationH
             return Box::new(|_, _| {});
         };
         Box::new(move |notifier, _| {
+            let _span = tracing::trace_span!("notification", method = N::METHOD).entered();
             if let Err(err) = N::run_with_snapshot(snapshot, notifier, params) {
                 tracing::error!("An error occurred while running {id}: {err}");
                 show_err_msg!("Ruff encountered a problem. Check the logs for more details.");
@@ -174,7 +183,7 @@ fn respond<Req>(
     Req: traits::RequestHandler,
 {
     if let Err(err) = &result {
-        tracing::error!("An error occurred with result ID {id}: {err}");
+        tracing::error!("An error occurred with request ID {id}: {err}");
         show_err_msg!("Ruff encountered a problem. Check the logs for more details.");
     }
     if let Err(err) = responder.respond(id, result) {

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

@@ -1,5 +1,5 @@
 use lsp_types::notification::DidOpenTextDocument;
-use lsp_types::DidOpenTextDocumentParams;
+use lsp_types::{DidOpenTextDocumentParams, TextDocumentItem};
 
 use red_knot_project::watch::ChangeEvent;
 use ruff_db::Db;
@@ -22,14 +22,22 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
         session: &mut Session,
         _notifier: Notifier,
         _requester: &mut Requester,
-        params: DidOpenTextDocumentParams,
+        DidOpenTextDocumentParams {
+            text_document:
+                TextDocumentItem {
+                    uri,
+                    text,
+                    version,
+                    language_id,
+                },
+        }: DidOpenTextDocumentParams,
     ) -> Result<()> {
-        let Ok(path) = url_to_any_system_path(&params.text_document.uri) else {
+        let Ok(path) = url_to_any_system_path(&uri) else {
             return Ok(());
         };
 
-        let document = TextDocument::new(params.text_document.text, params.text_document.version);
-        session.open_text_document(params.text_document.uri, document);
+        let document = TextDocument::new(text, version).with_language_id(&language_id);
+        session.open_text_document(uri, document);
 
         match path {
             AnySystemPath::System(path) => {

crates/red_knot_server/src/server/connection.rs

@@ -91,19 +91,40 @@ impl Connection {
                 self.sender
                     .send(lsp::Response::new_ok(id.clone(), ()).into())?;
                 tracing::info!("Shutdown request received. Waiting for an exit notification...");
-                match self.receiver.recv_timeout(std::time::Duration::from_secs(30))? {
-                    lsp::Message::Notification(lsp::Notification { method, .. }) if method == lsp_types::notification::Exit::METHOD => {
-                        tracing::info!("Exit notification received. Server shutting down...");
-                        Ok(true)
-                    },
-                    message => anyhow::bail!("Server received unexpected message {message:?} while waiting for exit notification")
+
+                loop {
+                    match &self
+                        .receiver
+                        .recv_timeout(std::time::Duration::from_secs(30))?
+                    {
+                        lsp::Message::Notification(lsp::Notification { method, .. })
+                            if method == lsp_types::notification::Exit::METHOD =>
+                        {
+                            tracing::info!("Exit notification received. Server shutting down...");
+                            return Ok(true);
+                        }
+                        lsp::Message::Request(lsp::Request { id, method, .. }) => {
+                            tracing::warn!(
+                                "Server received unexpected request {method} ({id}) while waiting for exit notification",
+                            );
+                            self.sender.send(lsp::Message::Response(lsp::Response::new_err(
+                                id.clone(),
+                                lsp::ErrorCode::InvalidRequest as i32,
+                                "Server received unexpected request while waiting for exit notification".to_string(),
+                            )))?;
+                        }
+                        message => {
+                            tracing::warn!(
+                                "Server received unexpected message while waiting for exit notification: {message:?}"
+                            );
+                        }
+                    }
                 }
             }
             lsp::Message::Notification(lsp::Notification { method, .. })
                 if method == lsp_types::notification::Exit::METHOD =>
             {
-                tracing::error!("Server received an exit notification before a shutdown request was sent. Exiting...");
-                Ok(true)
+                anyhow::bail!("Server received an exit notification before a shutdown request was sent. Exiting...");
             }
             _ => Ok(false),
         }