[red-knot] Support untitled files in the server (#13044)

## Summary

This PR adds support for untitled files in the red knot server.

## Test Plan

https://github.com/user-attachments/assets/57fa5db6-e1ad-4694-ae5f-c47a21eaa82b

crates/red_knot_server/src/server/api.rs

@@ -1,6 +1,9 @@
-use crate::{server::schedule::Task, session::Session, system::url_to_system_path};
 use lsp_server as server;
 
+use crate::server::schedule::Task;
+use crate::session::Session;
+use crate::system::{url_to_any_system_path, AnySystemPath};
+
 mod diagnostics;
 mod notifications;
 mod requests;
@@ -82,12 +85,17 @@ fn background_request_task<'a, R: traits::BackgroundDocumentRequestHandler>(
     Ok(Task::background(schedule, move |session: &Session| {
         let url = R::document_url(&params).into_owned();
 
-        let Ok(path) = url_to_system_path(&url) else {
+        let Ok(path) = url_to_any_system_path(&url) else {
             return Box::new(|_, _| {});
         };
-        let db = match session.workspace_db_for_path(path.as_std_path()) {
-            Some(db) => db.snapshot(),
-            None => session.default_workspace_db().snapshot(),
+        let db = match path {
+            AnySystemPath::System(path) => {
+                match session.workspace_db_for_path(path.as_std_path()) {
+                    Some(db) => db.snapshot(),
+                    None => session.default_workspace_db().snapshot(),
+                }
+            }
+            AnySystemPath::SystemVirtual(_) => session.default_workspace_db().snapshot(),
         };
 
         let Some(snapshot) = session.take_snapshot(url) else {

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

@@ -9,7 +9,7 @@ use crate::server::api::LSPResult;
 use crate::server::client::{Notifier, Requester};
 use crate::server::Result;
 use crate::session::Session;
-use crate::system::url_to_system_path;
+use crate::system::{url_to_any_system_path, AnySystemPath};
 
 pub(crate) struct DidChangeTextDocumentHandler;
 
@@ -24,7 +24,7 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
         _requester: &mut Requester,
         params: DidChangeTextDocumentParams,
     ) -> Result<()> {
-        let Ok(path) = url_to_system_path(&params.text_document.uri) else {
+        let Ok(path) = url_to_any_system_path(&params.text_document.uri) else {
             return Ok(());
         };
 
@@ -34,11 +34,19 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
             .update_text_document(&key, params.content_changes, params.text_document.version)
             .with_failure_code(ErrorCode::InternalError)?;
 
-        let db = match session.workspace_db_for_path_mut(path.as_std_path()) {
-            Some(db) => db,
-            None => session.default_workspace_db_mut(),
-        };
-        db.apply_changes(vec![ChangeEvent::file_content_changed(path)], None);
+        match path {
+            AnySystemPath::System(path) => {
+                let db = match session.workspace_db_for_path_mut(path.as_std_path()) {
+                    Some(db) => db,
+                    None => session.default_workspace_db_mut(),
+                };
+                db.apply_changes(vec![ChangeEvent::file_content_changed(path)], None);
+            }
+            AnySystemPath::SystemVirtual(virtual_path) => {
+                let db = session.default_workspace_db_mut();
+                db.apply_changes(vec![ChangeEvent::ChangedVirtual(virtual_path)], None);
+            }
+        }
 
         // TODO(dhruvmanila): Publish diagnostics if the client doesnt support pull diagnostics
 

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

@@ -1,6 +1,7 @@
 use lsp_server::ErrorCode;
 use lsp_types::notification::DidCloseTextDocument;
 use lsp_types::DidCloseTextDocumentParams;
+use red_knot_workspace::watch::ChangeEvent;
 
 use crate::server::api::diagnostics::clear_diagnostics;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
@@ -8,7 +9,7 @@ use crate::server::api::LSPResult;
 use crate::server::client::{Notifier, Requester};
 use crate::server::Result;
 use crate::session::Session;
-use crate::system::url_to_system_path;
+use crate::system::{url_to_any_system_path, AnySystemPath};
 
 pub(crate) struct DidCloseTextDocumentHandler;
 
@@ -23,7 +24,7 @@ impl SyncNotificationHandler for DidCloseTextDocumentHandler {
         _requester: &mut Requester,
         params: DidCloseTextDocumentParams,
     ) -> Result<()> {
-        let Ok(_path) = url_to_system_path(&params.text_document.uri) else {
+        let Ok(path) = url_to_any_system_path(&params.text_document.uri) else {
             return Ok(());
         };
 
@@ -32,6 +33,11 @@ impl SyncNotificationHandler for DidCloseTextDocumentHandler {
             .close_document(&key)
             .with_failure_code(ErrorCode::InternalError)?;
 
+        if let AnySystemPath::SystemVirtual(virtual_path) = path {
+            let db = session.default_workspace_db_mut();
+            db.apply_changes(vec![ChangeEvent::DeletedVirtual(virtual_path)], None);
+        }
+
         clear_diagnostics(key.url(), &notifier)?;
 
         Ok(())

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

@@ -1,12 +1,14 @@
 use lsp_types::notification::DidCloseNotebookDocument;
 use lsp_types::DidCloseNotebookDocumentParams;
 
+use red_knot_workspace::watch::ChangeEvent;
+
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::server::api::LSPResult;
 use crate::server::client::{Notifier, Requester};
 use crate::server::Result;
 use crate::session::Session;
-use crate::system::url_to_system_path;
+use crate::system::{url_to_any_system_path, AnySystemPath};
 
 pub(crate) struct DidCloseNotebookHandler;
 
@@ -21,7 +23,7 @@ impl SyncNotificationHandler for DidCloseNotebookHandler {
         _requester: &mut Requester,
         params: DidCloseNotebookDocumentParams,
     ) -> Result<()> {
-        let Ok(_path) = url_to_system_path(&params.notebook_document.uri) else {
+        let Ok(path) = url_to_any_system_path(&params.notebook_document.uri) else {
             return Ok(());
         };
 
@@ -30,6 +32,11 @@ impl SyncNotificationHandler for DidCloseNotebookHandler {
             .close_document(&key)
             .with_failure_code(lsp_server::ErrorCode::InternalError)?;
 
+        if let AnySystemPath::SystemVirtual(virtual_path) = path {
+            let db = session.default_workspace_db_mut();
+            db.apply_changes(vec![ChangeEvent::DeletedVirtual(virtual_path)], None);
+        }
+
         Ok(())
     }
 }

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

@@ -2,12 +2,13 @@ use lsp_types::notification::DidOpenTextDocument;
 use lsp_types::DidOpenTextDocumentParams;
 
 use red_knot_workspace::watch::ChangeEvent;
+use ruff_db::Db;
 
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
 use crate::server::client::{Notifier, Requester};
 use crate::server::Result;
 use crate::session::Session;
-use crate::system::url_to_system_path;
+use crate::system::{url_to_any_system_path, AnySystemPath};
 use crate::TextDocument;
 
 pub(crate) struct DidOpenTextDocumentHandler;
@@ -23,18 +24,26 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
         _requester: &mut Requester,
         params: DidOpenTextDocumentParams,
     ) -> Result<()> {
-        let Ok(path) = url_to_system_path(&params.text_document.uri) else {
+        let Ok(path) = url_to_any_system_path(&params.text_document.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 db = match session.workspace_db_for_path_mut(path.as_std_path()) {
-            Some(db) => db,
-            None => session.default_workspace_db_mut(),
-        };
-        db.apply_changes(vec![ChangeEvent::Opened(path)], None);
+        match path {
+            AnySystemPath::System(path) => {
+                let db = match session.workspace_db_for_path_mut(path.as_std_path()) {
+                    Some(db) => db,
+                    None => session.default_workspace_db_mut(),
+                };
+                db.apply_changes(vec![ChangeEvent::Opened(path)], None);
+            }
+            AnySystemPath::SystemVirtual(virtual_path) => {
+                let db = session.default_workspace_db_mut();
+                db.files().virtual_file(db, &virtual_path);
+            }
+        }
 
         // TODO(dhruvmanila): Publish diagnostics if the client doesn't support pull diagnostics
 

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

@@ -3,6 +3,7 @@ use lsp_types::notification::DidOpenNotebookDocument;
 use lsp_types::DidOpenNotebookDocumentParams;
 
 use red_knot_workspace::watch::ChangeEvent;
+use ruff_db::Db;
 
 use crate::edit::NotebookDocument;
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
@@ -10,7 +11,7 @@ use crate::server::api::LSPResult;
 use crate::server::client::{Notifier, Requester};
 use crate::server::Result;
 use crate::session::Session;
-use crate::system::url_to_system_path;
+use crate::system::{url_to_any_system_path, AnySystemPath};
 
 pub(crate) struct DidOpenNotebookHandler;
 
@@ -25,7 +26,7 @@ impl SyncNotificationHandler for DidOpenNotebookHandler {
         _requester: &mut Requester,
         params: DidOpenNotebookDocumentParams,
     ) -> Result<()> {
-        let Ok(path) = url_to_system_path(&params.notebook_document.uri) else {
+        let Ok(path) = url_to_any_system_path(&params.notebook_document.uri) else {
             return Ok(());
         };
 
@@ -38,11 +39,19 @@ impl SyncNotificationHandler for DidOpenNotebookHandler {
         .with_failure_code(ErrorCode::InternalError)?;
         session.open_notebook_document(params.notebook_document.uri.clone(), notebook);
 
-        let db = match session.workspace_db_for_path_mut(path.as_std_path()) {
-            Some(db) => db,
-            None => session.default_workspace_db_mut(),
-        };
-        db.apply_changes(vec![ChangeEvent::Opened(path)], None);
+        match path {
+            AnySystemPath::System(path) => {
+                let db = match session.workspace_db_for_path_mut(path.as_std_path()) {
+                    Some(db) => db,
+                    None => session.default_workspace_db_mut(),
+                };
+                db.apply_changes(vec![ChangeEvent::Opened(path)], None);
+            }
+            AnySystemPath::SystemVirtual(virtual_path) => {
+                let db = session.default_workspace_db_mut();
+                db.files().virtual_file(db, &virtual_path);
+            }
+        }
 
         // TODO(dhruvmanila): Publish diagnostics if the client doesn't support pull diagnostics
 

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

@@ -72,7 +72,7 @@ fn to_lsp_diagnostic(message: &str) -> Diagnostic {
     let words = message.split(':').collect::<Vec<_>>();
 
     let (range, message) = match words.as_slice() {
-        [_filename, line, column, message] => {
+        [_, _, line, column, message] | [_, line, column, message] => {
             let line = line.parse::<u32>().unwrap_or_default().saturating_sub(1);
             let column = column.parse::<u32>().unwrap_or_default();
             (

crates/red_knot_server/src/session.rs

@@ -12,9 +12,10 @@ use red_knot_workspace::db::RootDatabase;
 use red_knot_workspace::workspace::WorkspaceMetadata;
 use ruff_db::files::{system_path_to_file, File};
 use ruff_db::system::SystemPath;
+use ruff_db::Db;
 
 use crate::edit::{DocumentKey, DocumentVersion, NotebookDocument};
-use crate::system::{url_to_system_path, LSPSystem};
+use crate::system::{url_to_any_system_path, AnySystemPath, LSPSystem};
 use crate::{PositionEncoding, TextDocument};
 
 pub(crate) use self::capabilities::ResolvedClientCapabilities;
@@ -246,6 +247,7 @@ impl Drop for MutIndexGuard<'_> {
 
 /// An immutable snapshot of `Session` that references
 /// a specific document.
+#[derive(Debug)]
 pub struct DocumentSnapshot {
     resolved_client_capabilities: Arc<ResolvedClientCapabilities>,
     document_ref: index::DocumentQuery,
@@ -266,7 +268,12 @@ impl DocumentSnapshot {
     }
 
     pub(crate) fn file(&self, db: &RootDatabase) -> Option<File> {
-        let path = url_to_system_path(self.document_ref.file_url()).ok()?;
-        system_path_to_file(db, path).ok()
+        match url_to_any_system_path(self.document_ref.file_url()).ok()? {
+            AnySystemPath::System(path) => system_path_to_file(db, path).ok(),
+            AnySystemPath::SystemVirtual(virtual_path) => db
+                .files()
+                .try_virtual_file(&virtual_path)
+                .map(|virtual_file| virtual_file.file()),
+        }
     }
 }

crates/red_knot_server/src/system.rs

@@ -8,27 +8,40 @@ use ruff_db::file_revision::FileRevision;
 use ruff_db::system::walk_directory::WalkDirectoryBuilder;
 use ruff_db::system::{
     DirectoryEntry, FileType, Metadata, OsSystem, Result, System, SystemPath, SystemPathBuf,
-    SystemVirtualPath,
+    SystemVirtualPath, SystemVirtualPathBuf,
 };
 use ruff_notebook::{Notebook, NotebookError};
 
 use crate::session::index::Index;
 use crate::DocumentQuery;
 
-/// Converts the given [`Url`] to a [`SystemPathBuf`].
+/// Converts the given [`Url`] to an [`AnySystemPath`].
+///
+/// If the URL scheme is `file`, then the path is converted to a [`SystemPathBuf`]. Otherwise, the
+/// URL is converted to a [`SystemVirtualPathBuf`].
 ///
 /// This fails in the following cases:
-/// * The URL scheme is not `file`.
 /// * The URL cannot be converted to a file path (refer to [`Url::to_file_path`]).
 /// * If the URL is not a valid UTF-8 string.
-pub(crate) fn url_to_system_path(url: &Url) -> std::result::Result<SystemPathBuf, ()> {
+pub(crate) fn url_to_any_system_path(url: &Url) -> std::result::Result<AnySystemPath, ()> {
     if url.scheme() == "file" {
-        Ok(SystemPathBuf::from_path_buf(url.to_file_path()?).map_err(|_| ())?)
+        Ok(AnySystemPath::System(
+            SystemPathBuf::from_path_buf(url.to_file_path()?).map_err(|_| ())?,
+        ))
     } else {
-        Err(())
+        Ok(AnySystemPath::SystemVirtual(
+            SystemVirtualPath::new(url.as_str()).to_path_buf(),
+        ))
     }
 }
 
+/// Represents either a [`SystemPath`] or a [`SystemVirtualPath`].
+#[derive(Debug)]
+pub(crate) enum AnySystemPath {
+    System(SystemPathBuf),
+    SystemVirtual(SystemVirtualPathBuf),
+}
+
 #[derive(Debug)]
 pub(crate) struct LSPSystem {
     /// A read-only copy of the index where the server stores all the open documents and settings.

crates/red_knot_workspace/src/db/changes.rs

@@ -50,7 +50,7 @@ impl RootDatabase {
         };
 
         for change in changes {
-            if let Some(path) = change.path() {
+            if let Some(path) = change.system_path() {
                 if matches!(
                     path.file_name(),
                     Some(".gitignore" | ".ignore" | "ruff.toml" | ".ruff.toml" | "pyproject.toml")
@@ -131,6 +131,17 @@ impl RootDatabase {
                     }
                 }
 
+                watch::ChangeEvent::CreatedVirtual(path)
+                | watch::ChangeEvent::ChangedVirtual(path) => {
+                    File::sync_virtual_path(self, &path);
+                }
+
+                watch::ChangeEvent::DeletedVirtual(path) => {
+                    if let Some(virtual_file) = self.files().try_virtual_file(&path) {
+                        virtual_file.close(self);
+                    }
+                }
+
                 watch::ChangeEvent::Rescan => {
                     workspace_change = true;
                     Files::sync_all(self);

crates/red_knot_workspace/src/watch.rs

@@ -1,4 +1,4 @@
-use ruff_db::system::{SystemPath, SystemPathBuf};
+use ruff_db::system::{SystemPath, SystemPathBuf, SystemVirtualPathBuf};
 pub use watcher::{directory_watcher, EventHandler, Watcher};
 pub use workspace_watcher::WorkspaceWatcher;
 
@@ -41,6 +41,15 @@ pub enum ChangeEvent {
         kind: DeletedKind,
     },
 
+    /// A new virtual path was created.
+    CreatedVirtual(SystemVirtualPathBuf),
+
+    /// The content of a virtual path was changed.
+    ChangedVirtual(SystemVirtualPathBuf),
+
+    /// A virtual path was deleted.
+    DeletedVirtual(SystemVirtualPathBuf),
+
     /// The file watcher failed to observe some changes and now is out of sync with the file system.
     ///
     /// This can happen if many files are changed at once. The consumer should rescan all files to catch up
@@ -60,16 +69,16 @@ impl ChangeEvent {
     }
 
     pub fn file_name(&self) -> Option<&str> {
-        self.path().and_then(|path| path.file_name())
+        self.system_path().and_then(|path| path.file_name())
     }
 
-    pub fn path(&self) -> Option<&SystemPath> {
+    pub fn system_path(&self) -> Option<&SystemPath> {
         match self {
             ChangeEvent::Opened(path)
             | ChangeEvent::Created { path, .. }
             | ChangeEvent::Changed { path, .. }
             | ChangeEvent::Deleted { path, .. } => Some(path),
-            ChangeEvent::Rescan => None,
+            _ => None,
         }
     }
 }

crates/ruff_db/src/files.rs

@@ -318,6 +318,9 @@ impl File {
             }
             FilePath::Vendored(vendored) => db.vendored().read_to_string(vendored),
             FilePath::SystemVirtual(system_virtual) => {
+                // Add a dependency on the revision to ensure the operation gets re-executed when the file changes.
+                let _ = self.revision(db);
+
                 db.system().read_virtual_path_to_string(system_virtual)
             }
         }
@@ -342,6 +345,9 @@ impl File {
                 "Reading a notebook from the vendored file system is not supported.",
             ))),
             FilePath::SystemVirtual(system_virtual) => {
+                // Add a dependency on the revision to ensure the operation gets re-executed when the file changes.
+                let _ = self.revision(db);
+
                 db.system().read_virtual_path_to_notebook(system_virtual)
             }
         }