[ty] Add virtual files to the only project database (#19322)

## Summary

Previously, the virtual files were being added to the default database
that's present on the session. This is wrong because the default
database is for any files that don't belong to any project i.e., they're
outside of any projects managed by the server. Virtual files are neither
part of the project nor it is outside the projects. This was not the
intention as in the initial version, virtual files were being added to
the only project database managed by the server.

This PR fixes this by reverting back to the original behavior where
virtual files will be added to the only project database present. When
support for multiple workspace and project is added, this will require
updating (https://github.com/astral-sh/ty/issues/794).

This is required for #19264 because workspace diagnostics doesn't check
the default project database yet. Ideally, the default db should be
checked as well.

The implementation of this PR means that virtual files are now being
included for workspace diagnostics but it doesn't work completely e.g.,
if I save an untitled file the diagnostics disappears but it doesn't
appear back for the (now) saved file on disk as shown in the following
video demonstration:



https://github.com/user-attachments/assets/123e8d20-1e95-4c7d-b7eb-eb65be8c476e

crates/ty_server/src/server/api.rs

@@ -236,14 +236,7 @@ where
             });
         };
 
-        let db = match &path {
+        let db = session.project_db(&path).clone();
-            AnySystemPath::System(path) => match session.project_db_for_path(path) {
-                Some(db) => db.clone(),
-                None => session.default_project_db().clone(),
-            },
-            AnySystemPath::SystemVirtual(_) => session.default_project_db().clone(),
-        };
-
         let snapshot = session.take_document_snapshot(url);
 
         Box::new(move |client| {

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

@@ -82,7 +82,7 @@ pub(super) fn publish_diagnostics(session: &Session, key: &DocumentKey, client:
         }
     };
 
-    let db = session.project_db_or_default(key.path());
+    let db = session.project_db(key.path());
 
     let Some(diagnostics) = compute_diagnostics(db, &snapshot) else {
         return;

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

@@ -41,8 +41,8 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
             .with_failure_code(ErrorCode::InternalError)?;
 
         let changes = match key.path() {
-            AnySystemPath::System(path) => {
+            AnySystemPath::System(system_path) => {
-                vec![ChangeEvent::file_content_changed(path.clone())]
+                vec![ChangeEvent::file_content_changed(system_path.clone())]
             }
             AnySystemPath::SystemVirtual(virtual_path) => {
                 vec![ChangeEvent::ChangedVirtual(virtual_path.clone())]

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

@@ -44,12 +44,14 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
         let document = TextDocument::new(text, version).with_language_id(&language_id);
         session.open_text_document(key.path(), document);
 
-        match key.path() {
+        let path = key.path();
+
+        match path {
             AnySystemPath::System(system_path) => {
-                session.apply_changes(key.path(), vec![ChangeEvent::Opened(system_path.clone())]);
+                session.apply_changes(path, vec![ChangeEvent::Opened(system_path.clone())]);
             }
             AnySystemPath::SystemVirtual(virtual_path) => {
-                let db = session.default_project_db_mut();
+                let db = session.project_db_mut(path);
                 db.files().virtual_file(db, virtual_path);
             }
         }

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

@@ -43,7 +43,7 @@ impl SyncNotificationHandler for DidOpenNotebookHandler {
                 session.apply_changes(&path, vec![ChangeEvent::Opened(system_path.clone())]);
             }
             AnySystemPath::SystemVirtual(virtual_path) => {
-                let db = session.default_project_db_mut();
+                let db = session.project_db_mut(&path);
                 db.files().virtual_file(db, virtual_path);
             }
         }

crates/ty_server/src/session.rs

@@ -163,16 +163,47 @@ impl Session {
         &self.workspaces
     }
 
-    // TODO(dhruvmanila): Ideally, we should have a single method for `workspace_db_for_path_mut`
+    /// Returns a reference to the project's [`ProjectDatabase`] in which the given `path` belongs.
-    // and `default_workspace_db_mut` but the borrow checker doesn't allow that.
+    ///
-    // https://github.com/astral-sh/ruff/pull/13041#discussion_r1726725437
+    /// If the path is a system path, it will return the project database that is closest to the
+    /// given path, or the default project if no project is found for the path.
+    ///
+    /// If the path is a virtual path, it will return the first project database in the session.
+    pub(crate) fn project_db(&self, path: &AnySystemPath) -> &ProjectDatabase {
+        match path {
+            AnySystemPath::System(system_path) => self
+                .project_db_for_path(system_path)
+                .unwrap_or_else(|| self.default_project.get(self.index.as_ref())),
+            AnySystemPath::SystemVirtual(_virtual_path) => {
+                // TODO: Currently, ty only supports single workspace but we need to figure out
+                // which project should this virtual path belong to when there are multiple
+                // projects: https://github.com/astral-sh/ty/issues/794
+                self.projects.iter().next().map(|(_, db)| db).unwrap()
+            }
+        }
+    }
 
-    /// Returns a reference to the project's [`ProjectDatabase`] corresponding to the given path,
+    /// Returns a mutable reference to the project's [`ProjectDatabase`] in which the given `path`
-    /// or the default project if no project is found for the path.
+    /// belongs.
-    pub(crate) fn project_db_or_default(&self, path: &AnySystemPath) -> &ProjectDatabase {
+    ///
-        path.as_system()
+    /// Refer to [`project_db`] for more details on how the project is selected.
-            .and_then(|path| self.project_db_for_path(path))
+    ///
-            .unwrap_or_else(|| self.default_project_db())
+    /// [`project_db`]: Session::project_db
+    pub(crate) fn project_db_mut(&mut self, path: &AnySystemPath) -> &mut ProjectDatabase {
+        match path {
+            AnySystemPath::System(system_path) => self
+                .projects
+                .range_mut(..=system_path.to_path_buf())
+                .next_back()
+                .map(|(_, db)| db)
+                .unwrap_or_else(|| self.default_project.get_mut(self.index.as_ref())),
+            AnySystemPath::SystemVirtual(_virtual_path) => {
+                // TODO: Currently, ty only supports single workspace but we need to figure out
+                // which project should this virtual path belong to when there are multiple
+                // projects: https://github.com/astral-sh/ty/issues/794
+                self.projects.iter_mut().next().map(|(_, db)| db).unwrap()
+            }
+        }
     }
 
     /// Returns a reference to the project's [`ProjectDatabase`] corresponding to the given path, if
@@ -187,18 +218,6 @@ impl Session {
             .map(|(_, db)| db)
     }
 
-    /// Returns a mutable reference to the project [`ProjectDatabase`] corresponding to the given
-    /// path, if any.
-    pub(crate) fn project_db_for_path_mut(
-        &mut self,
-        path: impl AsRef<SystemPath>,
-    ) -> Option<&mut ProjectDatabase> {
-        self.projects
-            .range_mut(..=path.as_ref().to_path_buf())
-            .next_back()
-            .map(|(_, db)| db)
-    }
-
     pub(crate) fn apply_changes(
         &mut self,
         path: &AnySystemPath,
@@ -212,25 +231,8 @@ impl Session {
                 .cloned()
         });
 
-        let db = match path {
+        self.project_db_mut(path)
-            AnySystemPath::System(path) => match self.project_db_for_path_mut(path) {
+            .apply_changes(changes, overrides.as_ref())
-                Some(db) => db,
-                None => self.default_project_db_mut(),
-            },
-            AnySystemPath::SystemVirtual(_) => self.default_project_db_mut(),
-        };
-
-        db.apply_changes(changes, overrides.as_ref())
-    }
-
-    /// Returns a reference to the default project [`ProjectDatabase`].
-    pub(crate) fn default_project_db(&self) -> &ProjectDatabase {
-        self.default_project.get(self.index.as_ref())
-    }
-
-    /// Returns a mutable reference to the default project [`ProjectDatabase`].
-    pub(crate) fn default_project_db_mut(&mut self) -> &mut ProjectDatabase {
-        self.default_project.get_mut(self.index.as_ref())
     }
 
     /// Returns a mutable iterator over all project databases that have been initialized to this point.