diff --git a/crates/ruff_db/src/system/path.rs b/crates/ruff_db/src/system/path.rs
index 71a92fb4c8f67..a387ae54f62ed 100644
--- a/crates/ruff_db/src/system/path.rs
+++ b/crates/ruff_db/src/system/path.rs
@@ -723,10 +723,11 @@ impl ruff_cache::CacheKey for SystemPathBuf {
 
 /// A slice of a virtual path on [`System`](super::System) (akin to [`str`]).
 #[repr(transparent)]
+#[derive(Eq, PartialEq, Hash, PartialOrd, Ord)]
 pub struct SystemVirtualPath(str);
 
 impl SystemVirtualPath {
-    pub fn new(path: &str) -> &SystemVirtualPath {
+    pub const fn new(path: &str) -> &SystemVirtualPath {
         // SAFETY: SystemVirtualPath is marked as #[repr(transparent)] so the conversion from a
         // *const str to a *const SystemVirtualPath is valid.
         unsafe { &*(path as *const str as *const SystemVirtualPath) }
@@ -767,8 +768,8 @@ pub struct SystemVirtualPathBuf(String);
 
 impl SystemVirtualPathBuf {
     #[inline]
-    pub fn as_path(&self) -> &SystemVirtualPath {
-        SystemVirtualPath::new(&self.0)
+    pub const fn as_path(&self) -> &SystemVirtualPath {
+        SystemVirtualPath::new(self.0.as_str())
     }
 }
 
@@ -852,6 +853,12 @@ impl ruff_cache::CacheKey for SystemVirtualPathBuf {
     }
 }
 
+impl Borrow<SystemVirtualPath> for SystemVirtualPathBuf {
+    fn borrow(&self) -> &SystemVirtualPath {
+        self.as_path()
+    }
+}
+
 /// Deduplicates identical paths and removes nested paths.
 ///
 /// # Examples
diff --git a/crates/ty_server/src/document.rs b/crates/ty_server/src/document.rs
index fff51d2f493a1..e2c582475ba2e 100644
--- a/crates/ty_server/src/document.rs
+++ b/crates/ty_server/src/document.rs
@@ -11,6 +11,7 @@ use lsp_types::{PositionEncodingKind, Url};
 use crate::system::AnySystemPath;
 pub use notebook::NotebookDocument;
 pub(crate) use range::{FileRangeExt, PositionExt, RangeExt, TextSizeExt, ToRangeExt};
+use ruff_db::system::{SystemPathBuf, SystemVirtualPath};
 pub(crate) use text_document::DocumentVersion;
 pub use text_document::TextDocument;
 
@@ -41,39 +42,75 @@ impl From<PositionEncoding> for ruff_source_file::PositionEncoding {
 
 /// A unique document ID, derived from a URL passed as part of an LSP request.
 /// This document ID can point to either be a standalone Python file, a full notebook, or a cell within a notebook.
-#[derive(Clone, Debug)]
-pub(crate) enum DocumentKey {
-    Notebook(AnySystemPath),
-    NotebookCell {
-        cell_url: Url,
-        notebook_path: AnySystemPath,
-    },
-    Text(AnySystemPath),
+///
+/// The `DocumentKey` is very similar to `AnySystemPath`. The important distinction is that
+/// ty doesn't know about individual notebook cells, instead, ty operates on full notebook documents.
+/// ty also doesn't support resolving settings per cell, instead, settings are resolved per file or notebook.
+///
+/// Thus, the motivation of `DocumentKey` is to prevent accidental use of Cell keys for operations
+/// that expect to work on a file path level. That's what [`DocumentHandle::to_file_path`]
+/// is for, it returns a file path for any document, taking into account that these methods should
+/// return the notebook for cell documents and notebooks.
+#[derive(Clone, Debug, Hash, PartialEq, Eq)]
+pub(super) enum DocumentKey {
+    /// A URI using the `file` schema and maps to a valid path.
+    File(SystemPathBuf),
+
+    /// Any other URI.
+    ///
+    /// Used for Notebook-cells, URI's with non-`file` schemes, or invalid `file` URI's.
+    Opaque(String),
 }
 
 impl DocumentKey {
-    /// Returns the file path associated with the key.
-    pub(crate) fn path(&self) -> &AnySystemPath {
-        match self {
-            DocumentKey::Notebook(path) | DocumentKey::Text(path) => path,
-            DocumentKey::NotebookCell { notebook_path, .. } => notebook_path,
+    /// Converts the given [`Url`] to an [`DocumentKey`].
+    ///
+    /// If the URL scheme is `file`, then the path is converted to a [`SystemPathBuf`] unless
+    /// the url isn't a valid file path.
+    ///
+    /// In all other cases, the URL is kept as an opaque identifier ([`Self::Opaque`]).
+    pub(crate) fn from_url(url: &Url) -> Self {
+        if url.scheme() == "file" {
+            if let Ok(path) = url.to_file_path() {
+                Self::File(SystemPathBuf::from_path_buf(path).expect("URL to be valid UTF-8"))
+            } else {
+                tracing::warn!(
+                    "Treating `file:` url `{url}` as opaque URL as it isn't a valid file path"
+                );
+                Self::Opaque(url.to_string())
+            }
+        } else {
+            Self::Opaque(url.to_string())
         }
     }
 
-    pub(crate) fn from_path(path: AnySystemPath) -> Self {
-        // For text documents, we assume it's a text document unless it's a notebook file.
-        match path.extension() {
-            Some("ipynb") => Self::Notebook(path),
-            _ => Self::Text(path),
+    pub(crate) fn as_opaque(&self) -> Option<&str> {
+        match self {
+            Self::Opaque(uri) => Some(uri),
+            Self::File(_) => None,
         }
     }
 
-    /// Returns the URL for this document key. For notebook cells, returns the cell URL.
-    /// For other document types, converts the path to a URL.
-    pub(crate) fn to_url(&self) -> Option<Url> {
+    /// Returns the corresponding [`AnySystemPath`] for this document key.
+    ///
+    /// Note, calling this method on a `DocumentKey::Opaque` representing a cell document
+    /// will return a `SystemVirtualPath` corresponding to the cell URI but not the notebook file path.
+    /// That's most likely not what you want.
+    pub(super) fn to_file_path(&self) -> AnySystemPath {
         match self {
-            DocumentKey::NotebookCell { cell_url, .. } => Some(cell_url.clone()),
-            DocumentKey::Notebook(path) | DocumentKey::Text(path) => path.to_url(),
+            Self::File(path) => AnySystemPath::System(path.clone()),
+            Self::Opaque(uri) => {
+                AnySystemPath::SystemVirtual(SystemVirtualPath::new(uri).to_path_buf())
+            }
+        }
+    }
+}
+
+impl From<AnySystemPath> for DocumentKey {
+    fn from(value: AnySystemPath) -> Self {
+        match value {
+            AnySystemPath::System(system_path) => Self::File(system_path),
+            AnySystemPath::SystemVirtual(virtual_path) => Self::Opaque(virtual_path.to_string()),
         }
     }
 }
@@ -81,11 +118,8 @@ impl DocumentKey {
 impl std::fmt::Display for DocumentKey {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
-            Self::NotebookCell { cell_url, .. } => cell_url.fmt(f),
-            Self::Notebook(path) | Self::Text(path) => match path {
-                AnySystemPath::System(system_path) => system_path.fmt(f),
-                AnySystemPath::SystemVirtual(virtual_path) => virtual_path.fmt(f),
-            },
+            Self::File(path) => path.fmt(f),
+            Self::Opaque(uri) => uri.fmt(f),
         }
     }
 }
diff --git a/crates/ty_server/src/document/notebook.rs b/crates/ty_server/src/document/notebook.rs
index 2616cffd70611..d1e07648e2760 100644
--- a/crates/ty_server/src/document/notebook.rs
+++ b/crates/ty_server/src/document/notebook.rs
@@ -3,9 +3,8 @@ use lsp_types::NotebookCellKind;
 use ruff_notebook::CellMetadata;
 use rustc_hash::{FxBuildHasher, FxHashMap};
 
-use crate::{PositionEncoding, TextDocument};
-
 use super::DocumentVersion;
+use crate::{PositionEncoding, TextDocument};
 
 pub(super) type CellId = usize;
 
@@ -13,16 +12,25 @@ pub(super) type CellId = usize;
 /// contents are internally represented by [`TextDocument`]s.
 #[derive(Clone, Debug)]
 pub struct NotebookDocument {
+    url: lsp_types::Url,
     cells: Vec<NotebookCell>,
     metadata: ruff_notebook::RawNotebookMetadata,
     version: DocumentVersion,
     // Used to quickly find the index of a cell for a given URL.
-    cell_index: FxHashMap<lsp_types::Url, CellId>,
+    cell_index: FxHashMap<String, CellId>,
 }
 
 /// A single cell within a notebook, which has text contents represented as a `TextDocument`.
 #[derive(Clone, Debug)]
 struct NotebookCell {
+    /// The URL uniquely identifying the cell.
+    ///
+    /// > Cell text documents have a URI, but servers should not rely on any
+    /// > format for this URI, since it is up to the client on how it will
+    /// > create these URIs. The URIs must be unique across ALL notebook
+    /// > cells and can therefore be used to uniquely identify a notebook cell
+    /// >  or the cell’s text document.
+    /// > <https://microsoft.github.io/language-server-protocol/specifications/lsp/3.18/specification/#notebookDocument_synchronization>
     url: lsp_types::Url,
     kind: NotebookCellKind,
     document: TextDocument,
@@ -30,32 +38,45 @@ struct NotebookCell {
 
 impl NotebookDocument {
     pub fn new(
-        version: DocumentVersion,
+        url: lsp_types::Url,
+        notebook_version: DocumentVersion,
         cells: Vec<lsp_types::NotebookCell>,
         metadata: serde_json::Map<String, serde_json::Value>,
         cell_documents: Vec<lsp_types::TextDocumentItem>,
     ) -> crate::Result<Self> {
-        let mut cell_contents: FxHashMap<_, _> = cell_documents
-            .into_iter()
-            .map(|document| (document.uri, document.text))
-            .collect();
-
-        let cells: Vec<_> = cells
-            .into_iter()
-            .map(|cell| {
-                let contents = cell_contents.remove(&cell.document).unwrap_or_default();
-                NotebookCell::new(cell, contents, version)
-            })
-            .collect();
+        let mut cells: Vec<_> = cells.into_iter().map(NotebookCell::empty).collect();
+
+        let cell_index = Self::make_cell_index(&cells);
+
+        for cell_document in cell_documents {
+            let index = cell_index
+                .get(cell_document.uri.as_str())
+                .copied()
+                .ok_or_else(|| {
+                    anyhow::anyhow!(
+                        "Received content for cell `{}` that isn't present in the metadata",
+                        cell_document.uri
+                    )
+                })?;
+
+            cells[index].document =
+                TextDocument::new(cell_document.uri, cell_document.text, cell_document.version)
+                    .with_language_id(&cell_document.language_id);
+        }
 
         Ok(Self {
-            version,
-            cell_index: Self::make_cell_index(cells.as_slice()),
-            metadata: serde_json::from_value(serde_json::Value::Object(metadata))?,
+            url,
+            version: notebook_version,
+            cell_index,
             cells,
+            metadata: serde_json::from_value(serde_json::Value::Object(metadata))?,
         })
     }
 
+    pub(crate) fn url(&self) -> &lsp_types::Url {
+        &self.url
+    }
+
     /// Generates a pseudo-representation of a notebook that lacks per-cell metadata and contextual information
     /// but should still work with Ruff's linter.
     pub fn make_ruff_notebook(&self) -> ruff_notebook::Notebook {
@@ -127,7 +148,7 @@ impl NotebookDocument {
                 // First, delete the cells and remove them from the index.
                 if delete > 0 {
                     for cell in self.cells.drain(start..start + delete) {
-                        self.cell_index.remove(&cell.url);
+                        self.cell_index.remove(cell.url.as_str());
                         deleted_cells.insert(cell.url, cell.document);
                     }
                 }
@@ -150,7 +171,7 @@ impl NotebookDocument {
                 // Third, register the new cells in the index and update existing ones that came
                 // after the insertion.
                 for (index, cell) in self.cells.iter().enumerate().skip(start) {
-                    self.cell_index.insert(cell.url.clone(), index);
+                    self.cell_index.insert(cell.url.to_string(), index);
                 }
 
                 // Finally, update the text document that represents the cell with the actual
@@ -158,8 +179,9 @@ impl NotebookDocument {
                 // `cell_index` are updated before we start applying the changes to the cells.
                 if let Some(did_open) = structure.did_open {
                     for cell_text_document in did_open {
-                        if let Some(cell) = self.cell_by_uri_mut(&cell_text_document.uri) {
+                        if let Some(cell) = self.cell_by_uri_mut(cell_text_document.uri.as_str()) {
                             cell.document = TextDocument::new(
+                                cell_text_document.uri,
                                 cell_text_document.text,
                                 cell_text_document.version,
                             );
@@ -170,7 +192,7 @@ impl NotebookDocument {
 
             if let Some(cell_data) = data {
                 for cell in cell_data {
-                    if let Some(existing_cell) = self.cell_by_uri_mut(&cell.document) {
+                    if let Some(existing_cell) = self.cell_by_uri_mut(cell.document.as_str()) {
                         existing_cell.kind = cell.kind;
                     }
                 }
@@ -178,7 +200,7 @@ impl NotebookDocument {
 
             if let Some(content_changes) = text_content {
                 for content_change in content_changes {
-                    if let Some(cell) = self.cell_by_uri_mut(&content_change.document.uri) {
+                    if let Some(cell) = self.cell_by_uri_mut(content_change.document.uri.as_str()) {
                         cell.document
                             .apply_changes(content_change.changes, version, encoding);
                     }
@@ -204,7 +226,8 @@ impl NotebookDocument {
     }
 
     /// 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> {
+    #[expect(unused)]
+    pub(crate) fn cell_document_by_uri(&self, uri: &str) -> Option<&TextDocument> {
         self.cells
             .get(*self.cell_index.get(uri)?)
             .map(|cell| &cell.document)
@@ -215,29 +238,41 @@ impl NotebookDocument {
         self.cells.iter().map(|cell| &cell.url)
     }
 
-    fn cell_by_uri_mut(&mut self, uri: &lsp_types::Url) -> Option<&mut NotebookCell> {
+    fn cell_by_uri_mut(&mut self, uri: &str) -> Option<&mut NotebookCell> {
         self.cells.get_mut(*self.cell_index.get(uri)?)
     }
 
-    fn make_cell_index(cells: &[NotebookCell]) -> FxHashMap<lsp_types::Url, CellId> {
+    fn make_cell_index(cells: &[NotebookCell]) -> FxHashMap<String, CellId> {
         let mut index = FxHashMap::with_capacity_and_hasher(cells.len(), FxBuildHasher);
         for (i, cell) in cells.iter().enumerate() {
-            index.insert(cell.url.clone(), i);
+            index.insert(cell.url.to_string(), i);
         }
         index
     }
 }
 
 impl NotebookCell {
+    pub(crate) fn empty(cell: lsp_types::NotebookCell) -> Self {
+        Self {
+            kind: cell.kind,
+            document: TextDocument::new(
+                cell.document.clone(),
+                String::new(),
+                DocumentVersion::default(),
+            ),
+            url: cell.document,
+        }
+    }
+
     pub(crate) fn new(
         cell: lsp_types::NotebookCell,
         contents: String,
         version: DocumentVersion,
     ) -> Self {
         Self {
+            document: TextDocument::new(cell.document.clone(), contents, version),
             url: cell.document,
             kind: cell.kind,
-            document: TextDocument::new(contents, version),
         }
     }
 }
@@ -294,7 +329,14 @@ mod tests {
             }
         }
 
-        NotebookDocument::new(0, cells, serde_json::Map::default(), cell_documents).unwrap()
+        NotebookDocument::new(
+            lsp_types::Url::parse("file://test.ipynb").unwrap(),
+            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
diff --git a/crates/ty_server/src/document/text_document.rs b/crates/ty_server/src/document/text_document.rs
index e5d00ff0cf176..9898dd670b588 100644
--- a/crates/ty_server/src/document/text_document.rs
+++ b/crates/ty_server/src/document/text_document.rs
@@ -1,4 +1,4 @@
-use lsp_types::TextDocumentContentChangeEvent;
+use lsp_types::{TextDocumentContentChangeEvent, Url};
 use ruff_source_file::LineIndex;
 
 use crate::PositionEncoding;
@@ -11,6 +11,9 @@ pub(crate) type DocumentVersion = i32;
 /// with changes made by the user, including unsaved changes.
 #[derive(Debug, Clone)]
 pub struct TextDocument {
+    /// The URL as sent by the client
+    url: Url,
+
     /// The string contents of the document.
     contents: String,
     /// A computed line index for the document. This should always reflect
@@ -40,9 +43,10 @@ impl From<&str> for LanguageId {
 }
 
 impl TextDocument {
-    pub fn new(contents: String, version: DocumentVersion) -> Self {
+    pub fn new(url: Url, contents: String, version: DocumentVersion) -> Self {
         let index = LineIndex::from_source_text(&contents);
         Self {
+            url,
             contents,
             index,
             version,
@@ -60,6 +64,10 @@ impl TextDocument {
         self.contents
     }
 
+    pub(crate) fn url(&self) -> &Url {
+        &self.url
+    }
+
     pub fn contents(&self) -> &str {
         &self.contents
     }
@@ -154,11 +162,12 @@ impl TextDocument {
 #[cfg(test)]
 mod tests {
     use crate::{PositionEncoding, TextDocument};
-    use lsp_types::{Position, TextDocumentContentChangeEvent};
+    use lsp_types::{Position, TextDocumentContentChangeEvent, Url};
 
     #[test]
     fn redo_edit() {
         let mut document = TextDocument::new(
+            Url::parse("file:///test").unwrap(),
             r#""""
 测试comment
 一些测试内容
diff --git a/crates/ty_server/src/lib.rs b/crates/ty_server/src/lib.rs
index a56a95cb3884d..374c8421cfaa8 100644
--- a/crates/ty_server/src/lib.rs
+++ b/crates/ty_server/src/lib.rs
@@ -8,7 +8,7 @@ pub use crate::logging::{LogLevel, init_logging};
 pub use crate::server::{PartialWorkspaceProgress, PartialWorkspaceProgressParams, Server};
 pub use crate::session::{ClientOptions, DiagnosticMode};
 pub use document::{NotebookDocument, PositionEncoding, TextDocument};
-pub(crate) use session::{DocumentQuery, Session};
+pub(crate) use session::Session;
 
 mod capabilities;
 mod document;
diff --git a/crates/ty_server/src/server/api.rs b/crates/ty_server/src/server/api.rs
index 6fd1cde43a847..a56866791b967 100644
--- a/crates/ty_server/src/server/api.rs
+++ b/crates/ty_server/src/server/api.rs
@@ -1,6 +1,5 @@
 use crate::server::schedule::Task;
 use crate::session::Session;
-use crate::system::AnySystemPath;
 use anyhow::anyhow;
 use lsp_server as server;
 use lsp_server::RequestId;
@@ -208,7 +207,7 @@ where
 
         // SAFETY: The `snapshot` is safe to move across the unwind boundary because it is not used
         // after unwinding.
-        let snapshot = AssertUnwindSafe(session.take_session_snapshot());
+        let snapshot = AssertUnwindSafe(session.snapshot_session());
 
         Box::new(move |client| {
             let _span = tracing::debug_span!("request", %id, method = R::METHOD).entered();
@@ -253,10 +252,10 @@ where
             .cancellation_token(&id)
             .expect("request should have been tested for cancellation before scheduling");
 
-        let url = R::document_url(&params).into_owned();
+        let url = R::document_url(&params);
 
-        let Ok(path) = AnySystemPath::try_from_url(&url) else {
-            let reason = format!("URL `{url}` isn't a valid system path");
+        let Ok(document) = session.snapshot_document(&url) else {
+            let reason = format!("Document {url} is not open in the session");
             tracing::warn!(
                 "Ignoring request id={id} method={} because {reason}",
                 R::METHOD
@@ -274,8 +273,8 @@ where
             });
         };
 
+        let path = document.to_file_path();
         let db = session.project_db(&path).clone();
-        let snapshot = session.take_document_snapshot(url);
 
         Box::new(move |client| {
             let _span = tracing::debug_span!("request", %id, method = R::METHOD).entered();
@@ -294,7 +293,7 @@ where
             }
 
             if let Err(error) = ruff_db::panic::catch_unwind(|| {
-                R::handle_request(&id, &db, snapshot, client, params);
+                R::handle_request(&id, &db, document, client, params);
             }) {
                 panic_response::<R>(&id, client, &error, retry);
             }
@@ -371,7 +370,15 @@ where
     let (id, params) = cast_notification::<N>(req)?;
     Ok(Task::background(schedule, move |session: &Session| {
         let url = N::document_url(&params);
-        let snapshot = session.take_document_snapshot((*url).clone());
+        let Ok(snapshot) = session.snapshot_document(&url) else {
+            let reason = format!("Document {url} is not open in the session");
+            tracing::warn!(
+                "Ignoring notification id={id} method={} because {reason}",
+                N::METHOD
+            );
+            return Box::new(|_| {});
+        };
+
         Box::new(move |client| {
             let _span = tracing::debug_span!("notification", method = N::METHOD).entered();
 
diff --git a/crates/ty_server/src/server/api/diagnostics.rs b/crates/ty_server/src/server/api/diagnostics.rs
index d43b176a9b61c..7680dc1bad38b 100644
--- a/crates/ty_server/src/server/api/diagnostics.rs
+++ b/crates/ty_server/src/server/api/diagnostics.rs
@@ -13,16 +13,16 @@ use ruff_db::source::{line_index, source_text};
 use ruff_db::system::SystemPathBuf;
 use ty_project::{Db, ProjectDatabase};
 
-use crate::document::{DocumentKey, FileRangeExt, ToRangeExt};
+use crate::document::{FileRangeExt, ToRangeExt};
 use crate::session::DocumentSnapshot;
 use crate::session::client::Client;
 use crate::system::{AnySystemPath, file_to_url};
-use crate::{DocumentQuery, PositionEncoding, Session};
+use crate::{NotebookDocument, PositionEncoding, Session};
 
 pub(super) struct Diagnostics<'a> {
     items: Vec<ruff_db::diagnostic::Diagnostic>,
     encoding: PositionEncoding,
-    document: &'a DocumentQuery,
+    notebook: Option<&'a NotebookDocument>,
 }
 
 impl Diagnostics<'_> {
@@ -53,7 +53,7 @@ impl Diagnostics<'_> {
     }
 
     pub(super) fn to_lsp_diagnostics(&self, db: &ProjectDatabase) -> LspDiagnostics {
-        if let Some(notebook) = self.document.as_notebook() {
+        if let Some(notebook) = self.notebook {
             let mut cell_diagnostics: FxHashMap<Url, Vec<Diagnostic>> = FxHashMap::default();
 
             // Populates all relevant URLs with an empty diagnostic list. This ensures that documents
@@ -115,23 +115,18 @@ impl LspDiagnostics {
     }
 }
 
-/// Clears the diagnostics for the document identified by `key`.
+/// Clears the diagnostics for the document identified by `uri`.
 ///
 /// This is done by notifying the client with an empty list of diagnostics for the document.
 /// For notebook cells, this clears diagnostics for the specific cell.
 /// For other document types, this clears diagnostics for the main document.
-pub(super) fn clear_diagnostics(session: &Session, key: &DocumentKey, client: &Client) {
+pub(super) fn clear_diagnostics(session: &Session, uri: &lsp_types::Url, client: &Client) {
     if session.client_capabilities().supports_pull_diagnostics() {
         return;
     }
 
-    let Some(uri) = key.to_url() else {
-        // If we can't convert to URL, we can't clear diagnostics
-        return;
-    };
-
     client.send_notification::<PublishDiagnostics>(PublishDiagnosticsParams {
-        uri,
+        uri: uri.clone(),
         diagnostics: vec![],
         version: None,
     });
@@ -143,18 +138,12 @@ pub(super) fn clear_diagnostics(session: &Session, key: &DocumentKey, client: &C
 /// This function is a no-op if the client supports pull diagnostics.
 ///
 /// [publish diagnostics notification]: https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_publishDiagnostics
-pub(super) fn publish_diagnostics(session: &Session, key: &DocumentKey, client: &Client) {
+pub(super) fn publish_diagnostics(session: &Session, url: &lsp_types::Url, client: &Client) {
     if session.client_capabilities().supports_pull_diagnostics() {
         return;
     }
 
-    let Some(url) = key.to_url() else {
-        return;
-    };
-
-    let snapshot = session.take_document_snapshot(url.clone());
-
-    let document = match snapshot.document() {
+    let snapshot = match session.snapshot_document(url) {
         Ok(document) => document,
         Err(err) => {
             tracing::debug!("Failed to resolve document for URL `{}`: {}", url, err);
@@ -162,7 +151,7 @@ pub(super) fn publish_diagnostics(session: &Session, key: &DocumentKey, client:
         }
     };
 
-    let db = session.project_db(key.path());
+    let db = session.project_db(&snapshot.to_file_path());
 
     let Some(diagnostics) = compute_diagnostics(db, &snapshot) else {
         return;
@@ -173,13 +162,13 @@ pub(super) fn publish_diagnostics(session: &Session, key: &DocumentKey, client:
         client.send_notification::<PublishDiagnostics>(PublishDiagnosticsParams {
             uri,
             diagnostics,
-            version: Some(document.version()),
+            version: Some(snapshot.document().version()),
         });
     };
 
     match diagnostics.to_lsp_diagnostics(db) {
         LspDiagnostics::TextDocument(diagnostics) => {
-            publish_diagnostics_notification(url, diagnostics);
+            publish_diagnostics_notification(url.clone(), diagnostics);
         }
         LspDiagnostics::NotebookDocument(cell_diagnostics) => {
             for (cell_url, diagnostics) in cell_diagnostics {
@@ -264,16 +253,11 @@ pub(super) fn compute_diagnostics<'a>(
     db: &ProjectDatabase,
     snapshot: &'a DocumentSnapshot,
 ) -> Option<Diagnostics<'a>> {
-    let document = match snapshot.document() {
-        Ok(document) => document,
-        Err(err) => {
-            tracing::info!("Failed to resolve document for snapshot: {}", err);
-            return None;
-        }
-    };
-
-    let Some(file) = document.file(db) else {
-        tracing::info!("No file found for snapshot for `{}`", document.file_path());
+    let Some(file) = snapshot.to_file(db) else {
+        tracing::info!(
+            "No file found for snapshot for `{}`",
+            snapshot.to_file_path()
+        );
         return None;
     };
 
@@ -282,7 +266,7 @@ pub(super) fn compute_diagnostics<'a>(
     Some(Diagnostics {
         items: diagnostics,
         encoding: snapshot.encoding(),
-        document,
+        notebook: snapshot.notebook(),
     })
 }
 
diff --git a/crates/ty_server/src/server/api/notifications/did_change.rs b/crates/ty_server/src/server/api/notifications/did_change.rs
index 68f6f883e022b..3cb52c3daa9cc 100644
--- a/crates/ty_server/src/server/api/notifications/did_change.rs
+++ b/crates/ty_server/src/server/api/notifications/did_change.rs
@@ -28,19 +28,16 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
             content_changes,
         } = params;
 
-        let key = match session.key_from_url(uri) {
-            Ok(key) => key,
-            Err(uri) => {
-                tracing::debug!("Failed to create document key from URI: {}", uri);
-                return Ok(());
-            }
-        };
+        let document = session
+            .document_handle(&uri)
+            .with_failure_code(ErrorCode::InternalError)?;
 
-        session
-            .update_text_document(&key, content_changes, version)
+        document
+            .update_text_document(session, content_changes, version)
             .with_failure_code(ErrorCode::InternalError)?;
 
-        let changes = match key.path() {
+        let path = document.to_file_path();
+        let changes = match &*path {
             AnySystemPath::System(system_path) => {
                 vec![ChangeEvent::file_content_changed(system_path.clone())]
             }
@@ -49,9 +46,9 @@ impl SyncNotificationHandler for DidChangeTextDocumentHandler {
             }
         };
 
-        session.apply_changes(key.path(), changes);
+        session.apply_changes(&path, changes);
 
-        publish_diagnostics(session, &key, client);
+        publish_diagnostics(session, document.url(), client);
 
         Ok(())
     }
diff --git a/crates/ty_server/src/server/api/notifications/did_change_watched_files.rs b/crates/ty_server/src/server/api/notifications/did_change_watched_files.rs
index 21285f461f0c5..ce55100deeecd 100644
--- a/crates/ty_server/src/server/api/notifications/did_change_watched_files.rs
+++ b/crates/ty_server/src/server/api/notifications/did_change_watched_files.rs
@@ -1,3 +1,4 @@
+use crate::document::DocumentKey;
 use crate::server::Result;
 use crate::server::api::diagnostics::{publish_diagnostics, publish_settings_diagnostics};
 use crate::server::api::traits::{NotificationHandler, SyncNotificationHandler};
@@ -25,16 +26,8 @@ impl SyncNotificationHandler for DidChangeWatchedFiles {
         let mut events_by_db: FxHashMap<_, Vec<ChangeEvent>> = FxHashMap::default();
 
         for change in params.changes {
-            let path = match AnySystemPath::try_from_url(&change.uri) {
-                Ok(path) => path,
-                Err(err) => {
-                    tracing::warn!(
-                        "Failed to convert URI '{}` to system path: {err:?}",
-                        change.uri
-                    );
-                    continue;
-                }
-            };
+            let key = DocumentKey::from_url(&change.uri);
+            let path = key.to_file_path();
 
             let system_path = match path {
                 AnySystemPath::System(system) => system,
@@ -99,8 +92,8 @@ impl SyncNotificationHandler for DidChangeWatchedFiles {
                 |_, ()| {},
             );
         } else {
-            for key in session.text_document_keys() {
-                publish_diagnostics(session, &key, client);
+            for key in session.text_document_handles() {
+                publish_diagnostics(session, key.url(), client);
             }
         }
         // TODO: always publish diagnostics for notebook files (since they don't use pull diagnostics)
diff --git a/crates/ty_server/src/server/api/notifications/did_close.rs b/crates/ty_server/src/server/api/notifications/did_close.rs
index 60097df67b5c3..5c5747ee05575 100644
--- a/crates/ty_server/src/server/api/notifications/did_close.rs
+++ b/crates/ty_server/src/server/api/notifications/did_close.rs
@@ -27,22 +27,20 @@ impl SyncNotificationHandler for DidCloseTextDocumentHandler {
             text_document: TextDocumentIdentifier { uri },
         } = params;
 
-        let key = match session.key_from_url(uri) {
-            Ok(key) => key,
-            Err(uri) => {
-                tracing::debug!("Failed to create document key from URI: {}", uri);
-                return Ok(());
-            }
-        };
+        let document = session
+            .document_handle(&uri)
+            .with_failure_code(ErrorCode::InternalError)?;
+
+        let path = document.to_file_path().into_owned();
+        let url = document.url().clone();
 
-        session
-            .close_document(&key)
+        document
+            .close(session)
             .with_failure_code(ErrorCode::InternalError)?;
 
-        let path = key.path();
-        let db = session.project_db_mut(path);
+        let db = session.project_db_mut(&path);
 
-        match path {
+        match &path {
             AnySystemPath::System(system_path) => {
                 if let Some(file) = db.files().try_system(db, system_path) {
                     db.project().close_file(db, file);
@@ -65,7 +63,7 @@ impl SyncNotificationHandler for DidCloseTextDocumentHandler {
                         .diagnostic_mode()
                         .is_open_files_only()
                 {
-                    clear_diagnostics(session, &key, client);
+                    clear_diagnostics(session, &url, client);
                 }
             }
             AnySystemPath::SystemVirtual(virtual_path) => {
@@ -78,7 +76,7 @@ impl SyncNotificationHandler for DidCloseTextDocumentHandler {
 
                 // Always clear diagnostics for virtual files, as they don't really exist on disk
                 // which means closing them is like deleting the file.
-                clear_diagnostics(session, &key, client);
+                clear_diagnostics(session, &url, client);
             }
         }
 
diff --git a/crates/ty_server/src/server/api/notifications/did_close_notebook.rs b/crates/ty_server/src/server/api/notifications/did_close_notebook.rs
index f934f6832eb7b..9b03651496999 100644
--- a/crates/ty_server/src/server/api/notifications/did_close_notebook.rs
+++ b/crates/ty_server/src/server/api/notifications/did_close_notebook.rs
@@ -26,21 +26,19 @@ impl SyncNotificationHandler for DidCloseNotebookHandler {
             ..
         } = params;
 
-        let key = match session.key_from_url(uri) {
-            Ok(key) => key,
-            Err(uri) => {
-                tracing::debug!("Failed to create document key from URI: {}", uri);
-                return Ok(());
-            }
-        };
-
-        session
-            .close_document(&key)
+        let document = session
+            .document_handle(&uri)
             .with_failure_code(lsp_server::ErrorCode::InternalError)?;
 
-        if let AnySystemPath::SystemVirtual(virtual_path) = key.path() {
+        let path = document.to_file_path().into_owned();
+
+        document
+            .close(session)
+            .with_failure_code(lsp_server::ErrorCode::InternalError)?;
+
+        if let AnySystemPath::SystemVirtual(virtual_path) = &path {
             session.apply_changes(
-                key.path(),
+                &path,
                 vec![ChangeEvent::DeletedVirtual(virtual_path.clone())],
             );
         }
diff --git a/crates/ty_server/src/server/api/notifications/did_open.rs b/crates/ty_server/src/server/api/notifications/did_open.rs
index 5647bb2781025..b2561e9c6c32f 100644
--- a/crates/ty_server/src/server/api/notifications/did_open.rs
+++ b/crates/ty_server/src/server/api/notifications/did_open.rs
@@ -35,30 +35,23 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
                 },
         } = params;
 
-        let key = match session.key_from_url(uri) {
-            Ok(key) => key,
-            Err(uri) => {
-                tracing::debug!("Failed to create document key from URI: {}", uri);
-                return Ok(());
-            }
-        };
-
-        let document = TextDocument::new(text, version).with_language_id(&language_id);
-        session.open_text_document(key.path(), document);
+        let document = session.open_text_document(
+            TextDocument::new(uri, text, version).with_language_id(&language_id),
+        );
 
-        let path = key.path();
+        let path = document.to_file_path();
 
         // This is a "maybe" because the `File` might've not been interned yet i.e., the
         // `try_system` call will return `None` which doesn't mean that the file is new, it's just
         // that the server didn't need the file yet.
         let is_maybe_new_system_file = path.as_system().is_some_and(|system_path| {
-            let db = session.project_db(path);
+            let db = session.project_db(&path);
             db.files()
                 .try_system(db, system_path)
                 .is_none_or(|file| !file.exists(db))
         });
 
-        match path {
+        match &*path {
             AnySystemPath::System(system_path) => {
                 let event = if is_maybe_new_system_file {
                     ChangeEvent::Created {
@@ -68,22 +61,22 @@ impl SyncNotificationHandler for DidOpenTextDocumentHandler {
                 } else {
                     ChangeEvent::Opened(system_path.clone())
                 };
-                session.apply_changes(path, vec![event]);
+                session.apply_changes(&path, vec![event]);
 
-                let db = session.project_db_mut(path);
+                let db = session.project_db_mut(&path);
                 match system_path_to_file(db, system_path) {
                     Ok(file) => db.project().open_file(db, file),
                     Err(err) => tracing::warn!("Failed to open file {system_path}: {err}"),
                 }
             }
             AnySystemPath::SystemVirtual(virtual_path) => {
-                let db = session.project_db_mut(path);
+                let db = session.project_db_mut(&path);
                 let virtual_file = db.files().virtual_file(db, virtual_path);
                 db.project().open_file(db, virtual_file.file());
             }
         }
 
-        publish_diagnostics(session, &key, client);
+        publish_diagnostics(session, document.url(), client);
 
         Ok(())
     }
diff --git a/crates/ty_server/src/server/api/notifications/did_open_notebook.rs b/crates/ty_server/src/server/api/notifications/did_open_notebook.rs
index 201add9587f0b..b61f2aeef6d1f 100644
--- a/crates/ty_server/src/server/api/notifications/did_open_notebook.rs
+++ b/crates/ty_server/src/server/api/notifications/did_open_notebook.rs
@@ -25,20 +25,27 @@ impl SyncNotificationHandler for DidOpenNotebookHandler {
         _client: &Client,
         params: DidOpenNotebookDocumentParams,
     ) -> Result<()> {
-        let Ok(path) = AnySystemPath::try_from_url(&params.notebook_document.uri) else {
-            return Ok(());
-        };
+        let lsp_types::NotebookDocument {
+            version,
+            cells,
+            metadata,
+            uri: notebook_uri,
+            ..
+        } = params.notebook_document;
 
         let notebook = NotebookDocument::new(
-            params.notebook_document.version,
-            params.notebook_document.cells,
-            params.notebook_document.metadata.unwrap_or_default(),
+            notebook_uri,
+            version,
+            cells,
+            metadata.unwrap_or_default(),
             params.cell_text_documents,
         )
         .with_failure_code(ErrorCode::InternalError)?;
-        session.open_notebook_document(&path, notebook);
 
-        match &path {
+        let document = session.open_notebook_document(notebook);
+        let path = document.to_file_path();
+
+        match &*path {
             AnySystemPath::System(system_path) => {
                 session.apply_changes(&path, vec![ChangeEvent::Opened(system_path.clone())]);
             }
diff --git a/crates/ty_server/src/server/api/requests/completion.rs b/crates/ty_server/src/server/api/requests/completion.rs
index 35e1c3fd253c6..8f3ef353388e7 100644
--- a/crates/ty_server/src/server/api/requests/completion.rs
+++ b/crates/ty_server/src/server/api/requests/completion.rs
@@ -45,7 +45,7 @@ impl BackgroundDocumentRequestHandler for CompletionRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
+        let Some(file) = snapshot.to_file(db) else {
             return Ok(None);
         };
 
diff --git a/crates/ty_server/src/server/api/requests/doc_highlights.rs b/crates/ty_server/src/server/api/requests/doc_highlights.rs
index 9750bdc190dbc..b5b6d0d9ab727 100644
--- a/crates/ty_server/src/server/api/requests/doc_highlights.rs
+++ b/crates/ty_server/src/server/api/requests/doc_highlights.rs
@@ -37,7 +37,7 @@ impl BackgroundDocumentRequestHandler for DocumentHighlightRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
+        let Some(file) = snapshot.to_file(db) else {
             return Ok(None);
         };
 
diff --git a/crates/ty_server/src/server/api/requests/document_symbols.rs b/crates/ty_server/src/server/api/requests/document_symbols.rs
index 46c4c3eb2ead4..ea5ee312c6f34 100644
--- a/crates/ty_server/src/server/api/requests/document_symbols.rs
+++ b/crates/ty_server/src/server/api/requests/document_symbols.rs
@@ -39,7 +39,7 @@ impl BackgroundDocumentRequestHandler for DocumentSymbolRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
+        let Some(file) = snapshot.to_file(db) else {
             return Ok(None);
         };
 
diff --git a/crates/ty_server/src/server/api/requests/execute_command.rs b/crates/ty_server/src/server/api/requests/execute_command.rs
index 8a2fc52fd1c67..a51ece8598af9 100644
--- a/crates/ty_server/src/server/api/requests/execute_command.rs
+++ b/crates/ty_server/src/server/api/requests/execute_command.rs
@@ -52,7 +52,7 @@ fn debug_information(session: &Session) -> crate::Result<String> {
     writeln!(
         buffer,
         "Open text documents: {}",
-        session.text_document_keys().count()
+        session.text_document_handles().count()
     )?;
     writeln!(buffer)?;
 
diff --git a/crates/ty_server/src/server/api/requests/goto_declaration.rs b/crates/ty_server/src/server/api/requests/goto_declaration.rs
index 07444746f7098..1c16a74bc51ab 100644
--- a/crates/ty_server/src/server/api/requests/goto_declaration.rs
+++ b/crates/ty_server/src/server/api/requests/goto_declaration.rs
@@ -37,7 +37,7 @@ impl BackgroundDocumentRequestHandler for GotoDeclarationRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
+        let Some(file) = snapshot.to_file(db) else {
             return Ok(None);
         };
 
diff --git a/crates/ty_server/src/server/api/requests/goto_definition.rs b/crates/ty_server/src/server/api/requests/goto_definition.rs
index 793ae54bf12ec..bc334117781ef 100644
--- a/crates/ty_server/src/server/api/requests/goto_definition.rs
+++ b/crates/ty_server/src/server/api/requests/goto_definition.rs
@@ -37,7 +37,7 @@ impl BackgroundDocumentRequestHandler for GotoDefinitionRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
+        let Some(file) = snapshot.to_file(db) else {
             return Ok(None);
         };
 
diff --git a/crates/ty_server/src/server/api/requests/goto_references.rs b/crates/ty_server/src/server/api/requests/goto_references.rs
index 129afcecdc499..3afaf28b14756 100644
--- a/crates/ty_server/src/server/api/requests/goto_references.rs
+++ b/crates/ty_server/src/server/api/requests/goto_references.rs
@@ -37,7 +37,7 @@ impl BackgroundDocumentRequestHandler for ReferencesRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
+        let Some(file) = snapshot.to_file(db) else {
             return Ok(None);
         };
 
diff --git a/crates/ty_server/src/server/api/requests/goto_type_definition.rs b/crates/ty_server/src/server/api/requests/goto_type_definition.rs
index 5695c5a6abbf1..379defa3447c1 100644
--- a/crates/ty_server/src/server/api/requests/goto_type_definition.rs
+++ b/crates/ty_server/src/server/api/requests/goto_type_definition.rs
@@ -37,7 +37,7 @@ impl BackgroundDocumentRequestHandler for GotoTypeDefinitionRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
+        let Some(file) = snapshot.to_file(db) else {
             return Ok(None);
         };
 
diff --git a/crates/ty_server/src/server/api/requests/hover.rs b/crates/ty_server/src/server/api/requests/hover.rs
index be81eca4724d6..cc8f8e0dab2e4 100644
--- a/crates/ty_server/src/server/api/requests/hover.rs
+++ b/crates/ty_server/src/server/api/requests/hover.rs
@@ -37,7 +37,7 @@ impl BackgroundDocumentRequestHandler for HoverRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
+        let Some(file) = snapshot.to_file(db) else {
             return Ok(None);
         };
 
diff --git a/crates/ty_server/src/server/api/requests/inlay_hints.rs b/crates/ty_server/src/server/api/requests/inlay_hints.rs
index ec8464fc6b187..21eb1d09b626b 100644
--- a/crates/ty_server/src/server/api/requests/inlay_hints.rs
+++ b/crates/ty_server/src/server/api/requests/inlay_hints.rs
@@ -36,7 +36,7 @@ impl BackgroundDocumentRequestHandler for InlayHintRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
+        let Some(file) = snapshot.to_file(db) else {
             return Ok(None);
         };
 
diff --git a/crates/ty_server/src/server/api/requests/prepare_rename.rs b/crates/ty_server/src/server/api/requests/prepare_rename.rs
index 7f11961beeda2..a12541729de4c 100644
--- a/crates/ty_server/src/server/api/requests/prepare_rename.rs
+++ b/crates/ty_server/src/server/api/requests/prepare_rename.rs
@@ -37,7 +37,7 @@ impl BackgroundDocumentRequestHandler for PrepareRenameRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
+        let Some(file) = snapshot.to_file(db) else {
             return Ok(None);
         };
 
diff --git a/crates/ty_server/src/server/api/requests/rename.rs b/crates/ty_server/src/server/api/requests/rename.rs
index 117891ebba733..d434cb733eff0 100644
--- a/crates/ty_server/src/server/api/requests/rename.rs
+++ b/crates/ty_server/src/server/api/requests/rename.rs
@@ -38,7 +38,7 @@ impl BackgroundDocumentRequestHandler for RenameRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
+        let Some(file) = snapshot.to_file(db) else {
             return Ok(None);
         };
 
diff --git a/crates/ty_server/src/server/api/requests/selection_range.rs b/crates/ty_server/src/server/api/requests/selection_range.rs
index 684b230cd311d..516ea6aeda065 100644
--- a/crates/ty_server/src/server/api/requests/selection_range.rs
+++ b/crates/ty_server/src/server/api/requests/selection_range.rs
@@ -37,7 +37,7 @@ impl BackgroundDocumentRequestHandler for SelectionRangeRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
+        let Some(file) = snapshot.to_file(db) else {
             return Ok(None);
         };
 
diff --git a/crates/ty_server/src/server/api/requests/semantic_tokens.rs b/crates/ty_server/src/server/api/requests/semantic_tokens.rs
index 58f245d4aebae..adc6142189c55 100644
--- a/crates/ty_server/src/server/api/requests/semantic_tokens.rs
+++ b/crates/ty_server/src/server/api/requests/semantic_tokens.rs
@@ -33,7 +33,7 @@ impl BackgroundDocumentRequestHandler for SemanticTokensRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
+        let Some(file) = snapshot.to_file(db) else {
             return Ok(None);
         };
 
diff --git a/crates/ty_server/src/server/api/requests/semantic_tokens_range.rs b/crates/ty_server/src/server/api/requests/semantic_tokens_range.rs
index 61124052492ba..03193b32a6ea6 100644
--- a/crates/ty_server/src/server/api/requests/semantic_tokens_range.rs
+++ b/crates/ty_server/src/server/api/requests/semantic_tokens_range.rs
@@ -35,7 +35,7 @@ impl BackgroundDocumentRequestHandler for SemanticTokensRangeRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
+        let Some(file) = snapshot.to_file(db) else {
             return Ok(None);
         };
 
diff --git a/crates/ty_server/src/server/api/requests/signature_help.rs b/crates/ty_server/src/server/api/requests/signature_help.rs
index e9b9f160b6a8a..f9b20cccd92d5 100644
--- a/crates/ty_server/src/server/api/requests/signature_help.rs
+++ b/crates/ty_server/src/server/api/requests/signature_help.rs
@@ -39,7 +39,7 @@ impl BackgroundDocumentRequestHandler for SignatureHelpRequestHandler {
             return Ok(None);
         }
 
-        let Some(file) = snapshot.file(db) else {
+        let Some(file) = snapshot.to_file(db) else {
             return Ok(None);
         };
 
diff --git a/crates/ty_server/src/server/api/requests/workspace_diagnostic.rs b/crates/ty_server/src/server/api/requests/workspace_diagnostic.rs
index c990d4f4af259..2d3743611693f 100644
--- a/crates/ty_server/src/server/api/requests/workspace_diagnostic.rs
+++ b/crates/ty_server/src/server/api/requests/workspace_diagnostic.rs
@@ -1,4 +1,5 @@
 use crate::PositionEncoding;
+use crate::document::DocumentKey;
 use crate::server::api::diagnostics::{Diagnostics, to_lsp_diagnostic};
 use crate::server::api::traits::{
     BackgroundRequestHandler, RequestHandler, RetriableRequestHandler,
@@ -8,7 +9,7 @@ use crate::server::{Action, Result};
 use crate::session::client::Client;
 use crate::session::index::Index;
 use crate::session::{SessionSnapshot, SuspendedWorkspaceDiagnosticRequest};
-use crate::system::{AnySystemPath, file_to_url};
+use crate::system::file_to_url;
 use lsp_server::RequestId;
 use lsp_types::request::WorkspaceDiagnosticRequest;
 use lsp_types::{
@@ -317,7 +318,7 @@ struct ResponseWriter<'a> {
     // It's important that we use `AnySystemPath` over `Url` here because
     // `file_to_url` isn't guaranteed to return the exact same URL as the one provided
     // by the client.
-    previous_result_ids: FxHashMap<AnySystemPath, (Url, String)>,
+    previous_result_ids: FxHashMap<DocumentKey, (Url, String)>,
 }
 
 impl<'a> ResponseWriter<'a> {
@@ -346,12 +347,7 @@ impl<'a> ResponseWriter<'a> {
 
         let previous_result_ids = previous_result_ids
             .into_iter()
-            .filter_map(|prev| {
-                Some((
-                    AnySystemPath::try_from_url(&prev.uri).ok()?,
-                    (prev.uri, prev.value),
-                ))
-            })
+            .map(|prev| (DocumentKey::from_url(&prev.uri), (prev.uri, prev.value)))
             .collect();
 
         Self {
@@ -367,20 +363,16 @@ impl<'a> ResponseWriter<'a> {
             tracing::debug!("Failed to convert file path to URL at {}", file.path(db));
             return;
         };
-
+        let key = DocumentKey::from_url(&url);
         let version = self
             .index
-            .key_from_url(url.clone())
-            .ok()
-            .and_then(|key| self.index.make_document_ref(key).ok())
-            .map(|doc| i64::from(doc.version()));
+            .document_handle(&url)
+            .map(|doc| i64::from(doc.version()))
+            .ok();
 
         let result_id = Diagnostics::result_id_from_hash(diagnostics);
 
-        let previous_result_id = AnySystemPath::try_from_url(&url)
-            .ok()
-            .and_then(|path| self.previous_result_ids.remove(&path))
-            .map(|(_url, id)| id);
+        let previous_result_id = self.previous_result_ids.remove(&key).map(|(_url, id)| id);
 
         let report = match result_id {
             Some(new_id) if Some(&new_id) == previous_result_id.as_ref() => {
@@ -444,13 +436,12 @@ impl<'a> ResponseWriter<'a> {
 
         // Handle files that had diagnostics in previous request but no longer have any
         // Any remaining entries in previous_results are files that were fixed
-        for (previous_url, previous_result_id) in self.previous_result_ids.into_values() {
+        for (key, (previous_url, previous_result_id)) in self.previous_result_ids {
             // This file had diagnostics before but doesn't now, so we need to report it as having no diagnostics
             let version = self
                 .index
-                .key_from_url(previous_url.clone())
+                .document(&key)
                 .ok()
-                .and_then(|key| self.index.make_document_ref(key).ok())
                 .map(|doc| i64::from(doc.version()));
 
             let new_result_id = Diagnostics::result_id_from_hash(&[]);
diff --git a/crates/ty_server/src/session.rs b/crates/ty_server/src/session.rs
index 24ad0ef55e5c6..c5daec77e37f4 100644
--- a/crates/ty_server/src/session.rs
+++ b/crates/ty_server/src/session.rs
@@ -1,7 +1,7 @@
 //! Data model, state management, and configuration resolution.
 
 use anyhow::{Context, anyhow};
-use index::DocumentQueryError;
+use index::DocumentError;
 use lsp_server::{Message, RequestId};
 use lsp_types::notification::{DidChangeWatchedFiles, Exit, Notification};
 use lsp_types::request::{
@@ -15,8 +15,9 @@ use lsp_types::{
 };
 use options::GlobalOptions;
 use ruff_db::Db;
-use ruff_db::files::File;
+use ruff_db::files::{File, system_path_to_file};
 use ruff_db::system::{System, SystemPath, SystemPathBuf};
+use std::borrow::Cow;
 use std::collections::{BTreeMap, HashSet, VecDeque};
 use std::ops::{Deref, DerefMut};
 use std::panic::RefUnwindSafe;
@@ -26,7 +27,6 @@ use ty_project::metadata::Options;
 use ty_project::watch::ChangeEvent;
 use ty_project::{ChangeResult, CheckMode, Db as _, ProjectDatabase, ProjectMetadata};
 
-pub(crate) use self::index::DocumentQuery;
 pub(crate) use self::options::InitializationOptions;
 pub use self::options::{ClientOptions, DiagnosticMode};
 pub(crate) use self::settings::{GlobalSettings, WorkspaceSettings};
@@ -439,13 +439,6 @@ impl Session {
         self.projects.values_mut().chain(default_project)
     }
 
-    /// Returns the [`DocumentKey`] for the given URL.
-    ///
-    /// Refer to [`Index::key_from_url`] for more details.
-    pub(crate) fn key_from_url(&self, url: Url) -> Result<DocumentKey, Url> {
-        self.index().key_from_url(url)
-    }
-
     pub(crate) fn initialize_workspaces(
         &mut self,
         workspace_settings: Vec<(Url, ClientOptions)>,
@@ -819,25 +812,34 @@ impl Session {
     }
 
     /// Creates a document snapshot with the URL referencing the document to snapshot.
-    pub(crate) fn take_document_snapshot(&self, url: Url) -> DocumentSnapshot {
-        let key = self
-            .key_from_url(url)
-            .map_err(DocumentQueryError::InvalidUrl);
-        DocumentSnapshot {
+    pub(crate) fn snapshot_document(&self, url: &Url) -> Result<DocumentSnapshot, DocumentError> {
+        let index = self.index();
+        let document_handle = index.document_handle(url)?;
+
+        let notebook = if let Some(notebook_path) = &document_handle.notebook_path {
+            index
+                .notebook_arc(&DocumentKey::from(notebook_path.clone()))
+                .ok()
+        } else {
+            None
+        };
+
+        Ok(DocumentSnapshot {
             resolved_client_capabilities: self.resolved_client_capabilities,
             global_settings: self.global_settings.clone(),
-            workspace_settings: key
-                .as_ref()
-                .ok()
-                .and_then(|key| self.workspaces.settings_for_path(key.path().as_system()?))
+            workspace_settings: document_handle
+                .to_file_path()
+                .as_system()
+                .and_then(|path| self.workspaces.settings_for_path(path))
                 .unwrap_or_else(|| Arc::new(WorkspaceSettings::default())),
             position_encoding: self.position_encoding,
-            document_query_result: key.and_then(|key| self.index().make_document_ref(key)),
-        }
+            document: document_handle,
+            notebook,
+        })
     }
 
     /// Creates a snapshot of the current state of the [`Session`].
-    pub(crate) fn take_session_snapshot(&self) -> SessionSnapshot {
+    pub(crate) fn snapshot_session(&self) -> SessionSnapshot {
         SessionSnapshot {
             projects: self
                 .projects
@@ -855,56 +857,49 @@ impl Session {
     }
 
     /// Iterates over the document keys for all open text documents.
-    pub(super) fn text_document_keys(&self) -> impl Iterator<Item = DocumentKey> + '_ {
+    pub(super) fn text_document_handles(&self) -> impl Iterator<Item = DocumentHandle> + '_ {
         self.index()
-            .text_document_paths()
-            .map(|path| DocumentKey::Text(path.clone()))
+            .text_documents()
+            .map(|(key, document)| DocumentHandle {
+                key: key.clone(),
+                url: document.url().clone(),
+                version: document.version(),
+                // TODO: Set notebook path if text document is part of a notebook
+                notebook_path: None,
+            })
+    }
+
+    /// Returns a handle to the document specified by its URL.
+    ///
+    /// # Errors
+    ///
+    /// If the document is not found.
+    pub(crate) fn document_handle(
+        &self,
+        url: &lsp_types::Url,
+    ) -> Result<DocumentHandle, DocumentError> {
+        self.index().document_handle(url)
     }
 
     /// Registers a notebook document at the provided `path`.
     /// If a document is already open here, it will be overwritten.
-    pub(crate) fn open_notebook_document(
-        &mut self,
-        path: &AnySystemPath,
-        document: NotebookDocument,
-    ) {
-        self.index_mut().open_notebook_document(path, document);
+    ///
+    /// Returns a handle to the opened document.
+    pub(crate) fn open_notebook_document(&mut self, document: NotebookDocument) -> DocumentHandle {
+        let handle = self.index_mut().open_notebook_document(document);
         self.bump_revision();
+        handle
     }
 
     /// Registers a text document at the provided `path`.
     /// If a document is already open here, it will be overwritten.
-    pub(crate) fn open_text_document(&mut self, path: &AnySystemPath, document: TextDocument) {
-        self.index_mut().open_text_document(path, document);
-        self.bump_revision();
-    }
-
-    /// Updates a text document at the associated `key`.
     ///
-    /// The document key must point to a text document, or this will throw an error.
-    pub(crate) fn update_text_document(
-        &mut self,
-        key: &DocumentKey,
-        content_changes: Vec<TextDocumentContentChangeEvent>,
-        new_version: DocumentVersion,
-    ) -> crate::Result<()> {
-        let position_encoding = self.position_encoding;
-        self.index_mut().update_text_document(
-            key,
-            content_changes,
-            new_version,
-            position_encoding,
-        )?;
-        self.bump_revision();
-        Ok(())
-    }
+    /// Returns a handle to the opened document.
+    pub(crate) fn open_text_document(&mut self, document: TextDocument) -> DocumentHandle {
+        let handle = self.index_mut().open_text_document(document);
 
-    /// De-registers a document, specified by its key.
-    /// Calling this multiple times for the same document is a logic error.
-    pub(crate) fn close_document(&mut self, key: &DocumentKey) -> crate::Result<()> {
-        self.index_mut().close_document(key)?;
         self.bump_revision();
-        Ok(())
+        handle
     }
 
     /// Returns a reference to the index.
@@ -1003,7 +998,8 @@ pub(crate) struct DocumentSnapshot {
     global_settings: Arc<GlobalSettings>,
     workspace_settings: Arc<WorkspaceSettings>,
     position_encoding: PositionEncoding,
-    document_query_result: Result<DocumentQuery, DocumentQueryError>,
+    document: DocumentHandle,
+    notebook: Option<Arc<NotebookDocument>>,
 }
 
 impl DocumentSnapshot {
@@ -1028,27 +1024,28 @@ impl DocumentSnapshot {
     }
 
     /// Returns the result of the document query for this snapshot.
-    pub(crate) fn document(&self) -> Result<&DocumentQuery, &DocumentQueryError> {
-        self.document_query_result.as_ref()
+    pub(crate) fn document(&self) -> &DocumentHandle {
+        &self.document
     }
 
-    pub(crate) fn file(&self, db: &dyn Db) -> Option<File> {
-        let document = match self.document() {
-            Ok(document) => document,
-            Err(err) => {
-                tracing::debug!("Failed to resolve file: {}", err);
-                return None;
-            }
-        };
-        let file = document.file(db);
+    pub(crate) fn notebook(&self) -> Option<&NotebookDocument> {
+        self.notebook.as_deref()
+    }
+
+    pub(crate) fn to_file(&self, db: &dyn Db) -> Option<File> {
+        let file = self.document.to_file(db);
         if file.is_none() {
             tracing::debug!(
-                "Failed to resolve file: file not found for path `{}`",
-                document.file_path()
+                "Failed to resolve file: file not found for `{}`",
+                self.document.url()
             );
         }
         file
     }
+
+    pub(crate) fn to_file_path(&self) -> Cow<'_, AnySystemPath> {
+        self.document.to_file_path()
+    }
 }
 
 /// An immutable snapshot of the current state of [`Session`].
@@ -1320,3 +1317,90 @@ impl SuspendedWorkspaceDiagnosticRequest {
         None
     }
 }
+
+/// A handle to a document stored within [`Index`].
+///
+/// Allows identifying the document within the index but it also carries the URL used by the
+/// client to reference the document as well as the version of the document.
+///
+/// It also exposes methods to get the file-path of the corresponding ty-file.
+#[derive(Clone, Debug)]
+pub(crate) struct DocumentHandle {
+    /// The key that uniquely identifies this document in the index.
+    key: DocumentKey,
+    url: lsp_types::Url,
+    /// The path to the enclosing notebook file if this document is a notebook or a notebook cell.
+    notebook_path: Option<AnySystemPath>,
+    version: DocumentVersion,
+}
+
+impl DocumentHandle {
+    pub(crate) const fn version(&self) -> DocumentVersion {
+        self.version
+    }
+
+    /// The URL as used by the client to reference this document.
+    pub(crate) fn url(&self) -> &lsp_types::Url {
+        &self.url
+    }
+
+    /// The path to the enclosing file for this document.
+    ///
+    /// This is the path corresponding to the URL, except for notebook cells where the
+    /// path corresponds to the notebook file.
+    pub(crate) fn to_file_path(&self) -> Cow<'_, AnySystemPath> {
+        if let Some(path) = self.notebook_path.as_ref() {
+            Cow::Borrowed(path)
+        } else {
+            Cow::Owned(self.key.to_file_path())
+        }
+    }
+
+    /// Returns the salsa interned [`File`] for the document selected by this query.
+    ///
+    /// It returns [`None`] for the following cases:
+    /// - For virtual file, if it's not yet opened
+    /// - For regular file, if it does not exists or is a directory
+    pub(crate) fn to_file(&self, db: &dyn Db) -> Option<File> {
+        match &*self.to_file_path() {
+            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()),
+        }
+    }
+
+    pub(crate) fn update_text_document(
+        &self,
+        session: &mut Session,
+        content_changes: Vec<TextDocumentContentChangeEvent>,
+        new_version: DocumentVersion,
+    ) -> crate::Result<()> {
+        let position_encoding = session.position_encoding();
+        let mut index = session.index_mut();
+
+        let document_mut = index.document_mut(&self.key)?;
+
+        let Some(document) = document_mut.as_text_mut() else {
+            anyhow::bail!("Text document path does not point to a text document");
+        };
+
+        if content_changes.is_empty() {
+            document.update_version(new_version);
+            return Ok(());
+        }
+
+        document.apply_changes(content_changes, new_version, position_encoding);
+
+        Ok(())
+    }
+
+    /// De-registers a document, specified by its key.
+    /// Calling this multiple times for the same document is a logic error.
+    pub(crate) fn close(self, session: &mut Session) -> crate::Result<()> {
+        session.index_mut().close_document(&self.key)?;
+        session.bump_revision();
+        Ok(())
+    }
+}
diff --git a/crates/ty_server/src/session/index.rs b/crates/ty_server/src/session/index.rs
index 89d310f2ab8a4..95cc515a3585f 100644
--- a/crates/ty_server/src/session/index.rs
+++ b/crates/ty_server/src/session/index.rs
@@ -1,24 +1,24 @@
 use std::sync::Arc;
 
-use lsp_types::Url;
-use ruff_db::Db;
-use ruff_db::files::{File, system_path_to_file};
-use rustc_hash::FxHashMap;
-
+use crate::document::DocumentKey;
+use crate::session::DocumentHandle;
 use crate::{
     PositionEncoding, TextDocument,
-    document::{DocumentKey, DocumentVersion, NotebookDocument},
+    document::{DocumentVersion, NotebookDocument},
     system::AnySystemPath,
 };
 
+use ruff_db::system::SystemVirtualPath;
+use rustc_hash::FxHashMap;
+
 /// Stores and tracks all open documents in a session, along with their associated settings.
 #[derive(Debug)]
 pub(crate) struct Index {
     /// Maps all document file paths to the associated document controller
-    documents: FxHashMap<AnySystemPath, DocumentController>,
+    documents: FxHashMap<DocumentKey, Document>,
 
     /// Maps opaque cell URLs to a notebook path (document)
-    notebook_cells: FxHashMap<Url, AnySystemPath>,
+    notebook_cells: FxHashMap<String, AnySystemPath>,
 }
 
 impl Index {
@@ -29,68 +29,55 @@ impl Index {
         }
     }
 
-    pub(super) fn text_document_paths(&self) -> impl Iterator<Item = &AnySystemPath> + '_ {
-        self.documents
-            .iter()
-            .filter_map(|(path, doc)| doc.as_text().and(Some(path)))
-    }
-
-    #[expect(dead_code)]
-    pub(super) fn notebook_document_paths(&self) -> impl Iterator<Item = &AnySystemPath> + '_ {
-        self.documents
-            .iter()
-            .filter(|(_, doc)| doc.as_notebook().is_some())
-            .map(|(path, _)| path)
+    pub(super) fn text_documents(
+        &self,
+    ) -> impl Iterator<Item = (&DocumentKey, &TextDocument)> + '_ {
+        self.documents.iter().filter_map(|(key, doc)| {
+            let text_document = doc.as_text()?;
+            Some((key, text_document))
+        })
     }
 
-    pub(super) fn update_text_document(
-        &mut self,
-        key: &DocumentKey,
-        content_changes: Vec<lsp_types::TextDocumentContentChangeEvent>,
-        new_version: DocumentVersion,
-        encoding: PositionEncoding,
-    ) -> crate::Result<()> {
-        let controller = self.document_controller_for_key(key)?;
-        let Some(document) = controller.as_text_mut() else {
-            anyhow::bail!("Text document path does not point to a text document");
+    pub(crate) fn document_handle(
+        &self,
+        url: &lsp_types::Url,
+    ) -> Result<DocumentHandle, DocumentError> {
+        let key = DocumentKey::from_url(url);
+        let Some(document) = self.documents.get(&key) else {
+            return Err(DocumentError::NotFound(key));
         };
 
-        if content_changes.is_empty() {
-            document.update_version(new_version);
-            return Ok(());
+        if let Some(path) = key.as_opaque() {
+            if let Some(notebook_path) = self.notebook_cells.get(path) {
+                return Ok(DocumentHandle {
+                    key: key.clone(),
+                    notebook_path: Some(notebook_path.clone()),
+                    url: url.clone(),
+                    version: document.version(),
+                });
+            }
         }
 
-        document.apply_changes(content_changes, new_version, encoding);
-
-        Ok(())
+        Ok(DocumentHandle {
+            key: key.clone(),
+            notebook_path: None,
+            url: url.clone(),
+            version: document.version(),
+        })
     }
 
-    /// Returns the [`DocumentKey`] corresponding to the given URL.
-    ///
-    /// It returns [`Err`] with the original URL if it cannot be converted to a [`AnySystemPath`].
-    pub(crate) fn key_from_url(&self, url: Url) -> Result<DocumentKey, Url> {
-        if let Some(notebook_path) = self.notebook_cells.get(&url) {
-            Ok(DocumentKey::NotebookCell {
-                cell_url: url,
-                notebook_path: notebook_path.clone(),
-            })
-        } else {
-            let path = AnySystemPath::try_from_url(&url).map_err(|()| url)?;
-            if path
-                .extension()
-                .is_some_and(|ext| ext.eq_ignore_ascii_case("ipynb"))
-            {
-                Ok(DocumentKey::Notebook(path))
-            } else {
-                Ok(DocumentKey::Text(path))
-            }
-        }
+    #[expect(dead_code)]
+    pub(super) fn notebook_document_keys(&self) -> impl Iterator<Item = &DocumentKey> + '_ {
+        self.documents
+            .iter()
+            .filter(|(_, doc)| doc.as_notebook().is_some())
+            .map(|(key, _)| key)
     }
 
     #[expect(dead_code)]
     pub(super) fn update_notebook_document(
         &mut self,
-        key: &DocumentKey,
+        notebook_key: &DocumentKey,
         cells: Option<lsp_types::NotebookDocumentCellChange>,
         metadata: Option<serde_json::Map<String, serde_json::Value>>,
         new_version: DocumentVersion,
@@ -102,17 +89,16 @@ impl Index {
             ..
         }) = cells.as_ref().and_then(|cells| cells.structure.as_ref())
         {
-            let notebook_path = key.path().clone();
-
             for opened_cell in did_open {
+                let cell_path = SystemVirtualPath::new(opened_cell.uri.as_str());
                 self.notebook_cells
-                    .insert(opened_cell.uri.clone(), notebook_path.clone());
+                    .insert(cell_path.to_string(), notebook_key.to_file_path());
             }
             // deleted notebook cells are closed via textDocument/didClose - we don't close them here.
         }
 
-        let controller = self.document_controller_for_key(key)?;
-        let Some(notebook) = controller.as_notebook_mut() else {
+        let document = self.document_mut(notebook_key)?;
+        let Some(notebook) = document.as_notebook_mut() else {
             anyhow::bail!("Notebook document path does not point to a notebook document");
         };
 
@@ -123,44 +109,64 @@ impl Index {
     /// Create a document reference corresponding to the given document key.
     ///
     /// Returns an error if the document is not found or if the path cannot be converted to a URL.
-    pub(crate) fn make_document_ref(
-        &self,
-        key: DocumentKey,
-    ) -> Result<DocumentQuery, DocumentQueryError> {
-        let path = key.path();
-        let Some(controller) = self.documents.get(path) else {
-            return Err(DocumentQueryError::NotFound(key));
+    pub(crate) fn document(&self, key: &DocumentKey) -> Result<&Document, DocumentError> {
+        let Some(document) = self.documents.get(key) else {
+            return Err(DocumentError::NotFound(key.clone()));
         };
-        // TODO: The `to_url` conversion shouldn't be an error because the paths themselves are
-        // constructed from the URLs but the `Index` APIs don't maintain this invariant.
-        let (cell_url, file_path) = match key {
-            DocumentKey::NotebookCell {
-                cell_url,
-                notebook_path,
-            } => (Some(cell_url), notebook_path),
-            DocumentKey::Notebook(path) | DocumentKey::Text(path) => (None, path),
+
+        Ok(document)
+    }
+
+    pub(crate) fn notebook_arc(
+        &self,
+        key: &DocumentKey,
+    ) -> Result<Arc<NotebookDocument>, DocumentError> {
+        let Some(document) = self.documents.get(key) else {
+            return Err(DocumentError::NotFound(key.clone()));
         };
-        Ok(controller.make_ref(cell_url, file_path))
+
+        if let Document::Notebook(notebook) = document {
+            Ok(notebook.clone())
+        } else {
+            Err(DocumentError::NotFound(key.clone()))
+        }
     }
 
-    pub(super) fn open_text_document(&mut self, path: &AnySystemPath, document: TextDocument) {
-        self.documents
-            .insert(path.clone(), DocumentController::new_text(document));
+    pub(super) fn open_text_document(&mut self, document: TextDocument) -> DocumentHandle {
+        let key = DocumentKey::from_url(document.url());
+
+        // TODO: Fix file path for notebook cells
+        let handle = DocumentHandle {
+            key: key.clone(),
+            notebook_path: None,
+            url: document.url().clone(),
+            version: document.version(),
+        };
+
+        self.documents.insert(key, Document::new_text(document));
+
+        handle
     }
 
-    pub(super) fn open_notebook_document(
-        &mut self,
-        notebook_path: &AnySystemPath,
-        document: NotebookDocument,
-    ) {
+    pub(super) fn open_notebook_document(&mut self, document: NotebookDocument) -> DocumentHandle {
+        let notebook_key = DocumentKey::from_url(document.url());
+        let url = document.url().clone();
+        let version = document.version();
+
         for cell_url in document.cell_urls() {
             self.notebook_cells
-                .insert(cell_url.clone(), notebook_path.clone());
+                .insert(cell_url.to_string(), notebook_key.to_file_path());
+        }
+
+        self.documents
+            .insert(notebook_key.clone(), Document::new_notebook(document));
+
+        DocumentHandle {
+            notebook_path: Some(notebook_key.to_file_path()),
+            key: notebook_key,
+            url,
+            version,
         }
-        self.documents.insert(
-            notebook_path.clone(),
-            DocumentController::new_notebook(document),
-        );
     }
 
     pub(super) fn close_document(&mut self, key: &DocumentKey) -> crate::Result<()> {
@@ -169,27 +175,23 @@ impl Index {
         // is requested to be `closed` by VS Code after the notebook gets updated.
         // This is not documented in the LSP specification explicitly, and this assumption
         // may need revisiting in the future as we support more editors with notebook support.
-        if let DocumentKey::NotebookCell { cell_url, .. } = key {
-            if self.notebook_cells.remove(cell_url).is_none() {
-                tracing::warn!("Tried to remove a notebook cell that does not exist: {cell_url}");
-            }
-            return Ok(());
+        if let DocumentKey::Opaque(uri) = key {
+            self.notebook_cells.remove(uri);
         }
-        let path = key.path();
 
-        let Some(_) = self.documents.remove(path) else {
+        let Some(_) = self.documents.remove(key) else {
             anyhow::bail!("tried to close document that didn't exist at {key}")
         };
+
         Ok(())
     }
 
-    fn document_controller_for_key(
+    pub(super) fn document_mut(
         &mut self,
         key: &DocumentKey,
-    ) -> crate::Result<&mut DocumentController> {
-        let path = key.path();
-        let Some(controller) = self.documents.get_mut(path) else {
-            anyhow::bail!("Document controller not available at `{key}`");
+    ) -> Result<&mut Document, DocumentError> {
+        let Some(controller) = self.documents.get_mut(key) else {
+            return Err(DocumentError::NotFound(key.clone()));
         };
         Ok(controller)
     }
@@ -197,31 +199,24 @@ impl Index {
 
 /// A mutable handler to an underlying document.
 #[derive(Debug)]
-enum DocumentController {
+pub(crate) enum Document {
     Text(Arc<TextDocument>),
     Notebook(Arc<NotebookDocument>),
 }
 
-impl DocumentController {
-    fn new_text(document: TextDocument) -> Self {
+impl Document {
+    pub(super) fn new_text(document: TextDocument) -> Self {
         Self::Text(Arc::new(document))
     }
 
-    fn new_notebook(document: NotebookDocument) -> Self {
+    pub(super) fn new_notebook(document: NotebookDocument) -> Self {
         Self::Notebook(Arc::new(document))
     }
 
-    fn make_ref(&self, cell_url: Option<Url>, file_path: AnySystemPath) -> DocumentQuery {
-        match &self {
-            Self::Notebook(notebook) => DocumentQuery::Notebook {
-                cell_url,
-                file_path,
-                notebook: notebook.clone(),
-            },
-            Self::Text(document) => DocumentQuery::Text {
-                file_path,
-                document: document.clone(),
-            },
+    pub(crate) fn version(&self) -> DocumentVersion {
+        match self {
+            Self::Text(document) => document.version(),
+            Self::Notebook(notebook) => notebook.version(),
         }
     }
 
@@ -254,85 +249,8 @@ impl DocumentController {
     }
 }
 
-/// A read-only query to an open document.
-///
-/// This query can 'select' a text document, full notebook, or a specific notebook cell.
-/// It also includes document settings.
-#[derive(Debug, Clone)]
-pub(crate) enum DocumentQuery {
-    Text {
-        file_path: AnySystemPath,
-        document: Arc<TextDocument>,
-    },
-    Notebook {
-        /// The selected notebook cell, if it exists.
-        cell_url: Option<Url>,
-        /// The path to the notebook.
-        file_path: AnySystemPath,
-        notebook: Arc<NotebookDocument>,
-    },
-}
-
-impl DocumentQuery {
-    /// Attempts to access the underlying notebook document that this query is selecting.
-    pub(crate) fn as_notebook(&self) -> Option<&NotebookDocument> {
-        match self {
-            Self::Notebook { notebook, .. } => Some(notebook),
-            Self::Text { .. } => None,
-        }
-    }
-
-    /// Get the version of document selected by this query.
-    pub(crate) fn version(&self) -> DocumentVersion {
-        match self {
-            Self::Text { document, .. } => document.version(),
-            Self::Notebook { notebook, .. } => notebook.version(),
-        }
-    }
-
-    /// Get the system path for the document selected by this query.
-    pub(crate) fn file_path(&self) -> &AnySystemPath {
-        match self {
-            Self::Text { file_path, .. } | Self::Notebook { file_path, .. } => file_path,
-        }
-    }
-
-    /// Attempt to access the single inner text document selected by the query.
-    /// If this query is selecting an entire notebook document, this will return `None`.
-    #[expect(dead_code)]
-    pub(crate) fn as_single_document(&self) -> Option<&TextDocument> {
-        match self {
-            Self::Text { document, .. } => Some(document),
-            Self::Notebook {
-                notebook,
-                cell_url: cell_uri,
-                ..
-            } => cell_uri
-                .as_ref()
-                .and_then(|cell_uri| notebook.cell_document_by_uri(cell_uri)),
-        }
-    }
-
-    /// Returns the salsa interned [`File`] for the document selected by this query.
-    ///
-    /// It returns [`None`] for the following cases:
-    /// - For virtual file, if it's not yet opened
-    /// - For regular file, if it does not exists or is a directory
-    pub(crate) fn file(&self, db: &dyn Db) -> Option<File> {
-        match self.file_path() {
-            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()),
-        }
-    }
-}
-
 #[derive(Debug, Clone, thiserror::Error)]
-pub(crate) enum DocumentQueryError {
-    #[error("invalid URL: {0}")]
-    InvalidUrl(Url),
+pub(crate) enum DocumentError {
     #[error("document not found for key: {0}")]
     NotFound(DocumentKey),
 }
diff --git a/crates/ty_server/src/system.rs b/crates/ty_server/src/system.rs
index 323e4a68463c7..17b9bcbde665e 100644
--- a/crates/ty_server/src/system.rs
+++ b/crates/ty_server/src/system.rs
@@ -4,6 +4,8 @@ use std::fmt::Display;
 use std::panic::RefUnwindSafe;
 use std::sync::Arc;
 
+use crate::document::DocumentKey;
+use crate::session::index::{Document, Index};
 use lsp_types::Url;
 use ruff_db::file_revision::FileRevision;
 use ruff_db::files::{File, FilePath};
@@ -16,10 +18,6 @@ use ruff_notebook::{Notebook, NotebookError};
 use ty_ide::cached_vendored_path;
 use ty_python_semantic::Db;
 
-use crate::DocumentQuery;
-use crate::document::DocumentKey;
-use crate::session::index::Index;
-
 /// Returns a [`Url`] for the given [`File`].
 pub(crate) fn file_to_url(db: &dyn Db, file: File) -> Option<Url> {
     match file.path(db) {
@@ -41,26 +39,6 @@ pub(crate) enum AnySystemPath {
 }
 
 impl AnySystemPath {
-    /// 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 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 try_from_url(url: &Url) -> std::result::Result<Self, ()> {
-        if url.scheme() == "file" {
-            Ok(AnySystemPath::System(
-                SystemPathBuf::from_path_buf(url.to_file_path()?).map_err(|_| ())?,
-            ))
-        } else {
-            Ok(AnySystemPath::SystemVirtual(
-                SystemVirtualPath::new(url.as_str()).to_path_buf(),
-            ))
-        }
-    }
-
     pub(crate) const fn as_system(&self) -> Option<&SystemPathBuf> {
         match self {
             AnySystemPath::System(system_path_buf) => Some(system_path_buf),
@@ -68,21 +46,11 @@ impl AnySystemPath {
         }
     }
 
-    /// Returns the extension of the path, if any.
-    pub(crate) fn extension(&self) -> Option<&str> {
-        match self {
-            AnySystemPath::System(system_path) => system_path.extension(),
-            AnySystemPath::SystemVirtual(virtual_path) => virtual_path.extension(),
-        }
-    }
-
-    /// Converts the path to a URL.
-    pub(crate) fn to_url(&self) -> Option<Url> {
+    #[expect(unused)]
+    pub(crate) const fn as_virtual(&self) -> Option<&SystemVirtualPath> {
         match self {
-            AnySystemPath::System(system_path) => {
-                Url::from_file_path(system_path.as_std_path()).ok()
-            }
-            AnySystemPath::SystemVirtual(virtual_path) => Url::parse(virtual_path.as_str()).ok(),
+            AnySystemPath::SystemVirtual(path) => Some(path.as_path()),
+            AnySystemPath::System(_) => None,
         }
     }
 }
@@ -144,21 +112,17 @@ impl LSPSystem {
         self.index.as_ref().unwrap()
     }
 
-    fn make_document_ref(&self, path: AnySystemPath) -> Option<DocumentQuery> {
+    fn make_document_ref(&self, path: AnySystemPath) -> Option<&Document> {
         let index = self.index();
-        let key = DocumentKey::from_path(path);
-        index.make_document_ref(key).ok()
+        index.document(&DocumentKey::from(path)).ok()
     }
 
-    fn system_path_to_document_ref(&self, path: &SystemPath) -> Option<DocumentQuery> {
+    fn system_path_to_document_ref(&self, path: &SystemPath) -> Option<&Document> {
         let any_path = AnySystemPath::System(path.to_path_buf());
         self.make_document_ref(any_path)
     }
 
-    fn system_virtual_path_to_document_ref(
-        &self,
-        path: &SystemVirtualPath,
-    ) -> Option<DocumentQuery> {
+    fn system_virtual_path_to_document_ref(&self, path: &SystemVirtualPath) -> Option<&Document> {
         let any_path = AnySystemPath::SystemVirtual(path.to_path_buf());
         self.make_document_ref(any_path)
     }
@@ -170,7 +134,7 @@ impl System for LSPSystem {
 
         if let Some(document) = document {
             Ok(Metadata::new(
-                document_revision(&document),
+                document_revision(document),
                 None,
                 FileType::File,
             ))
@@ -191,7 +155,7 @@ impl System for LSPSystem {
         let document = self.system_path_to_document_ref(path);
 
         match document {
-            Some(DocumentQuery::Text { document, .. }) => Ok(document.contents().to_string()),
+            Some(Document::Text(document)) => Ok(document.contents().to_string()),
             _ => self.native_system.read_to_string(path),
         }
     }
@@ -200,10 +164,8 @@ impl System for LSPSystem {
         let document = self.system_path_to_document_ref(path);
 
         match document {
-            Some(DocumentQuery::Text { document, .. }) => {
-                Notebook::from_source_code(document.contents())
-            }
-            Some(DocumentQuery::Notebook { notebook, .. }) => Ok(notebook.make_ruff_notebook()),
+            Some(Document::Text(document)) => Notebook::from_source_code(document.contents()),
+            Some(Document::Notebook(notebook)) => Ok(notebook.make_ruff_notebook()),
             None => self.native_system.read_to_notebook(path),
         }
     }
@@ -213,7 +175,7 @@ impl System for LSPSystem {
             .system_virtual_path_to_document_ref(path)
             .ok_or_else(|| virtual_path_not_found(path))?;
 
-        if let DocumentQuery::Text { document, .. } = &document {
+        if let Document::Text(document) = &document {
             Ok(document.contents().to_string())
         } else {
             Err(not_a_text_document(path))
@@ -229,8 +191,8 @@ impl System for LSPSystem {
             .ok_or_else(|| virtual_path_not_found(path))?;
 
         match document {
-            DocumentQuery::Text { document, .. } => Notebook::from_source_code(document.contents()),
-            DocumentQuery::Notebook { notebook, .. } => Ok(notebook.make_ruff_notebook()),
+            Document::Text(document) => Notebook::from_source_code(document.contents()),
+            Document::Notebook(notebook) => Ok(notebook.make_ruff_notebook()),
         }
     }
 
@@ -307,7 +269,7 @@ fn virtual_path_not_found(path: impl Display) -> std::io::Error {
 }
 
 /// Helper function to get the [`FileRevision`] of the given document.
-fn document_revision(document: &DocumentQuery) -> FileRevision {
+fn document_revision(document: &Document) -> FileRevision {
     // The file revision is just an opaque number which doesn't have any significant meaning other
     // than that the file has changed if the revisions are different.
     #[expect(clippy::cast_sign_loss)]