refactor(lsp): determine file referrer for each document (#23867)

7 files changed, 300 insertions, 94 deletions

diff --git a/cli/lsp/analysis.rs b/cli/lsp/analysis.rs
index 133f5f0ac..9c7025781 100644
--- a/cli/lsp/analysis.rs
+++ b/cli/lsp/analysis.rs
@@ -248,6 +248,8 @@ impl<'a> TsResponseImportMapper<'a> {
       }
     }
 
+    let file_referrer = self.documents.get_file_referrer(referrer);
+
     if let Some(jsr_path) = specifier.as_str().strip_prefix(jsr_url().as_str())
     {
       let mut segments = jsr_path.split('/');
@@ -259,7 +261,11 @@ impl<'a> TsResponseImportMapper<'a> {
       let version = Version::parse_standard(segments.next()?).ok()?;
       let nv = PackageNv { name, version };
       let path = segments.collect::<Vec<_>>().join("/");
-      let export = self.resolver.jsr_lookup_export_for_path(&nv, &path)?;
+      let export = self.resolver.jsr_lookup_export_for_path(
+        &nv,
+        &path,
+        file_referrer.as_deref(),
+      )?;
       let sub_path = (export != ".").then_some(export);
       let mut req = None;
       req = req.or_else(|| {
@@ -281,7 +287,11 @@ impl<'a> TsResponseImportMapper<'a> {
         }
         None
       });
-      req = req.or_else(|| self.resolver.jsr_lookup_req_for_nv(&nv));
+      req = req.or_else(|| {
+        self
+          .resolver
+          .jsr_lookup_req_for_nv(&nv, file_referrer.as_deref())
+      });
       let spec_str = if let Some(req) = req {
         let req_ref = PackageReqReference { req, sub_path };
         JsrPackageReqReference::new(req_ref).to_string()
@@ -298,7 +308,10 @@ impl<'a> TsResponseImportMapper<'a> {
       return Some(spec_str);
     }
 
-    if let Some(npm_resolver) = self.resolver.maybe_managed_npm_resolver() {
+    if let Some(npm_resolver) = self
+      .resolver
+      .maybe_managed_npm_resolver(file_referrer.as_deref())
+    {
       if npm_resolver.in_npm_package(specifier) {
         if let Ok(Some(pkg_id)) =
           npm_resolver.resolve_pkg_id_from_specifier(specifier)
diff --git a/cli/lsp/completions.rs b/cli/lsp/completions.rs
index 3f63d2857..92e602ec6 100644
--- a/cli/lsp/completions.rs
+++ b/cli/lsp/completions.rs
@@ -157,6 +157,7 @@ pub async fn get_import_completions(
   maybe_import_map: Option<&ImportMap>,
 ) -> Option<lsp::CompletionResponse> {
   let document = documents.get(specifier)?;
+  let file_referrer = document.file_referrer();
   let (text, _, range) = document.get_maybe_dependency(position)?;
   let range = to_narrow_lsp_range(&document.text_info(), &range);
   if let Some(completion_list) = get_import_map_completions(
@@ -209,8 +210,8 @@ pub async fn get_import_completions(
       0
     };
     let maybe_list = module_registries
-      .get_completions(&text, offset, &range, |specifier| {
-        documents.exists(specifier)
+      .get_completions(&text, offset, &range, |s| {
+        documents.exists(s, file_referrer)
       })
       .await;
     let list = maybe_list.unwrap_or_else(|| lsp::CompletionList {
@@ -825,7 +826,7 @@ mod tests {
     for (specifier, source, version, language_id) in open_sources {
       let specifier =
         resolve_url(specifier).expect("failed to create specifier");
-      documents.open(specifier, *version, *language_id, (*source).into());
+      documents.open(specifier, *version, *language_id, (*source).into(), None);
     }
     for (specifier, source) in fs_sources {
       let specifier =
@@ -834,10 +835,9 @@ mod tests {
         .global()
         .set(&specifier, HashMap::default(), source.as_bytes())
         .expect("could not cache file");
-      assert!(
-        documents.get(&specifier).is_some(),
-        "source could not be setup"
-      );
+      let document =
+        documents.get_or_load(&specifier, &temp_dir.uri().join("$").unwrap());
+      assert!(document.is_some(), "source could not be setup");
     }
     documents
   }
diff --git a/cli/lsp/diagnostics.rs b/cli/lsp/diagnostics.rs
index 8472ad185..5c982b0f3 100644
--- a/cli/lsp/diagnostics.rs
+++ b/cli/lsp/diagnostics.rs
@@ -1294,6 +1294,7 @@ fn diagnose_resolution(
   resolution: &Resolution,
   is_dynamic: bool,
   maybe_assert_type: Option<&str>,
+  referrer_doc: &Document,
   import_map: Option<&ImportMap>,
 ) -> Vec<DenoDiagnostic> {
   fn check_redirect_diagnostic(
@@ -1327,13 +1328,21 @@ fn diagnose_resolution(
   match resolution {
     Resolution::Ok(resolved) => {
       let specifier = &resolved.specifier;
-      let managed_npm_resolver = snapshot.resolver.maybe_managed_npm_resolver();
-      for (_, headers) in snapshot.resolver.redirect_chain_headers(specifier) {
+      let managed_npm_resolver = snapshot
+        .resolver
+        .maybe_managed_npm_resolver(referrer_doc.file_referrer());
+      for (_, headers) in snapshot
+        .resolver
+        .redirect_chain_headers(specifier, referrer_doc.file_referrer())
+      {
         if let Some(message) = headers.get("x-deno-warning") {
           diagnostics.push(DenoDiagnostic::DenoWarn(message.clone()));
         }
       }
-      if let Some(doc) = snapshot.documents.get(specifier) {
+      if let Some(doc) = snapshot
+        .documents
+        .get_or_load(specifier, referrer_doc.specifier())
+      {
         if let Some(headers) = doc.maybe_headers() {
           if let Some(message) = headers.get("x-deno-warning") {
             diagnostics.push(DenoDiagnostic::DenoWarn(message.clone()));
@@ -1430,10 +1439,11 @@ fn diagnose_resolution(
 fn diagnose_dependency(
   diagnostics: &mut Vec<lsp::Diagnostic>,
   snapshot: &language_server::StateSnapshot,
-  referrer: &ModuleSpecifier,
+  referrer_doc: &Document,
   dependency_key: &str,
   dependency: &deno_graph::Dependency,
 ) {
+  let referrer = referrer_doc.specifier();
   if snapshot.resolver.in_node_modules(referrer) {
     return; // ignore, surface typescript errors instead
   }
@@ -1488,6 +1498,7 @@ fn diagnose_dependency(
       },
       dependency.is_dynamic,
       dependency.maybe_attribute_type.as_deref(),
+      referrer_doc,
       import_map.map(|i| i.as_ref()),
     )
     .iter()
@@ -1511,6 +1522,7 @@ fn diagnose_dependency(
         &dependency.maybe_type,
         dependency.is_dynamic,
         dependency.maybe_attribute_type.as_deref(),
+        referrer_doc,
         import_map.map(|i| i.as_ref()),
       )
       .iter()
@@ -1543,7 +1555,7 @@ fn generate_deno_diagnostics(
         diagnose_dependency(
           &mut diagnostics,
           snapshot,
-          specifier,
+          &document,
           dependency_key,
           dependency,
         );
@@ -1630,11 +1642,12 @@ mod tests {
         *version,
         *language_id,
         (*source).into(),
+        None,
       );
     }
     StateSnapshot {
       project_version: 0,
-      documents,
+      documents: Arc::new(documents),
       assets: Default::default(),
       config: Arc::new(config),
       resolver,
diff --git a/cli/lsp/documents.rs b/cli/lsp/documents.rs
index df15bff2a..ab053336f 100644
--- a/cli/lsp/documents.rs
+++ b/cli/lsp/documents.rs
@@ -266,6 +266,9 @@ pub struct Document {
   /// Contains the last-known-good set of dependencies from parsing the module.
   config: Arc<Config>,
   dependencies: Arc<IndexMap<String, deno_graph::Dependency>>,
+  // TODO(nayeemrmn): This is unused, use it for scope attribution for remote
+  // modules.
+  file_referrer: Option<ModuleSpecifier>,
   maybe_types_dependency: Option<Arc<deno_graph::TypesDependency>>,
   maybe_fs_version: Option<String>,
   line_index: Arc<LineIndex>,
@@ -295,6 +298,7 @@ impl Document {
     resolver: Arc<LspResolver>,
     config: Arc<Config>,
     cache: &Arc<LspCache>,
+    file_referrer: Option<ModuleSpecifier>,
   ) -> Arc<Self> {
     let text_info = SourceTextInfo::new(content);
     let media_type = resolve_media_type(
@@ -310,6 +314,7 @@ impl Document {
           text_info.clone(),
           maybe_headers.as_ref(),
           media_type,
+          file_referrer.as_ref(),
           &resolver,
         )
       } else {
@@ -329,6 +334,7 @@ impl Document {
     Arc::new(Self {
       config,
       dependencies,
+      file_referrer: file_referrer.filter(|_| specifier.scheme() != "file"),
       maybe_types_dependency,
       maybe_fs_version: calculate_fs_version(cache, &specifier),
       line_index,
@@ -370,6 +376,7 @@ impl Document {
         &self.specifier,
         &parsed_source_result,
         self.maybe_headers.as_ref(),
+        self.file_referrer.as_ref(),
         &resolver,
       )
       .ok();
@@ -384,8 +391,10 @@ impl Document {
       maybe_test_module_fut =
         get_maybe_test_module_fut(maybe_parsed_source.as_ref(), &config);
     } else {
-      let graph_resolver = resolver.as_graph_resolver();
-      let npm_resolver = resolver.as_graph_npm_resolver();
+      let graph_resolver =
+        resolver.as_graph_resolver(self.file_referrer.as_ref());
+      let npm_resolver =
+        resolver.as_graph_npm_resolver(self.file_referrer.as_ref());
       dependencies = Arc::new(
         self
           .dependencies
@@ -420,6 +429,7 @@ impl Document {
       config,
       // updated properties
       dependencies,
+      file_referrer: self.file_referrer.clone(),
       maybe_types_dependency,
       maybe_navigation_tree: Mutex::new(None),
       // maintain - this should all be copies/clones
@@ -475,6 +485,7 @@ impl Document {
         text_info.clone(),
         self.maybe_headers.as_ref(),
         media_type,
+        self.file_referrer.as_ref(),
         self.resolver.as_ref(),
       )
     } else {
@@ -499,6 +510,7 @@ impl Document {
     Ok(Arc::new(Self {
       config: self.config.clone(),
       specifier: self.specifier.clone(),
+      file_referrer: self.file_referrer.clone(),
       maybe_fs_version: self.maybe_fs_version.clone(),
       maybe_language_id: self.maybe_language_id,
       dependencies,
@@ -522,6 +534,7 @@ impl Document {
     Arc::new(Self {
       config: self.config.clone(),
       specifier: self.specifier.clone(),
+      file_referrer: self.file_referrer.clone(),
       maybe_fs_version: calculate_fs_version(cache, &self.specifier),
       maybe_language_id: self.maybe_language_id,
       dependencies: self.dependencies.clone(),
@@ -543,6 +556,7 @@ impl Document {
     Arc::new(Self {
       config: self.config.clone(),
       specifier: self.specifier.clone(),
+      file_referrer: self.file_referrer.clone(),
       maybe_fs_version: calculate_fs_version(cache, &self.specifier),
       maybe_language_id: self.maybe_language_id,
       dependencies: self.dependencies.clone(),
@@ -564,6 +578,10 @@ impl Document {
     &self.specifier
   }
 
+  pub fn file_referrer(&self) -> Option<&ModuleSpecifier> {
+    self.file_referrer.as_ref()
+  }
+
   pub fn content(&self) -> Arc<str> {
     self.text_info.text()
   }
@@ -729,6 +747,7 @@ impl FileSystemDocuments {
     resolver: &Arc<LspResolver>,
     config: &Arc<Config>,
     cache: &Arc<LspCache>,
+    file_referrer: Option<&ModuleSpecifier>,
   ) -> Option<Arc<Document>> {
     let new_fs_version = calculate_fs_version(cache, specifier);
     let old_doc = self.docs.get(specifier).map(|v| v.value().clone());
@@ -745,7 +764,7 @@ impl FileSystemDocuments {
     };
     if dirty {
       // attempt to update the file on the file system
-      self.refresh_document(specifier, resolver, config, cache)
+      self.refresh_document(specifier, resolver, config, cache, file_referrer)
     } else {
       old_doc
     }
@@ -759,6 +778,7 @@ impl FileSystemDocuments {
     resolver: &Arc<LspResolver>,
     config: &Arc<Config>,
     cache: &Arc<LspCache>,
+    file_referrer: Option<&ModuleSpecifier>,
   ) -> Option<Arc<Document>> {
     let doc = if specifier.scheme() == "file" {
       let path = specifier_to_file_path(specifier).ok()?;
@@ -774,6 +794,7 @@ impl FileSystemDocuments {
         resolver.clone(),
         config.clone(),
         cache,
+        file_referrer.cloned(),
       )
     } else if specifier.scheme() == "data" {
       let source = deno_graph::source::RawDataUrl::parse(specifier)
@@ -789,6 +810,7 @@ impl FileSystemDocuments {
         resolver.clone(),
         config.clone(),
         cache,
+        file_referrer.cloned(),
       )
     } else {
       let http_cache = cache.root_vendor_or_global();
@@ -818,6 +840,7 @@ impl FileSystemDocuments {
         resolver.clone(),
         config.clone(),
         cache,
+        file_referrer.cloned(),
       )
     };
     self.docs.insert(specifier.clone(), doc.clone());
@@ -882,6 +905,7 @@ impl Documents {
     version: i32,
     language_id: LanguageId,
     content: Arc<str>,
+    file_referrer: Option<ModuleSpecifier>,
   ) -> Arc<Document> {
     let document = Document::new(
       specifier.clone(),
@@ -895,6 +919,7 @@ impl Documents {
       self.resolver.clone(),
       self.config.clone(),
       &self.cache,
+      file_referrer,
     );
 
     self.file_system_docs.remove_document(&specifier);
@@ -964,6 +989,19 @@ impl Documents {
     self.file_system_docs.set_dirty(true);
   }
 
+  pub fn get_file_referrer<'a>(
+    &self,
+    specifier: &'a ModuleSpecifier,
+  ) -> Option<Cow<'a, ModuleSpecifier>> {
+    if specifier.scheme() == "file" {
+      Some(Cow::Borrowed(specifier))
+    } else {
+      self
+        .get(specifier)
+        .and_then(|d| d.file_referrer().cloned().map(Cow::Owned))
+    }
+  }
+
   /// Return `true` if the provided specifier can be resolved to a document,
   /// otherwise `false`.
   pub fn contains_import(
@@ -971,9 +1009,10 @@ impl Documents {
     specifier: &str,
     referrer: &ModuleSpecifier,
   ) -> bool {
+    let file_referrer = self.get_file_referrer(referrer);
     let maybe_specifier = self
       .resolver
-      .as_graph_resolver()
+      .as_graph_resolver(file_referrer.as_deref())
       .resolve(
         specifier,
         &deno_graph::Range {
@@ -985,7 +1024,7 @@ impl Documents {
       )
       .ok();
     if let Some(import_specifier) = maybe_specifier {
-      self.exists(&import_specifier)
+      self.exists(&import_specifier, file_referrer.as_deref())
     } else {
       false
     }
@@ -994,23 +1033,32 @@ impl Documents {
   pub fn resolve_document_specifier(
     &self,
     specifier: &ModuleSpecifier,
+    file_referrer: Option<&ModuleSpecifier>,
   ) -> Option<ModuleSpecifier> {
     let specifier = if let Ok(jsr_req_ref) =
       JsrPackageReqReference::from_specifier(specifier)
     {
-      Cow::Owned(self.resolver.jsr_to_registry_url(&jsr_req_ref)?)
+      Cow::Owned(
+        self
+          .resolver
+          .jsr_to_registry_url(&jsr_req_ref, file_referrer)?,
+      )
     } else {
       Cow::Borrowed(specifier)
     };
     if !DOCUMENT_SCHEMES.contains(&specifier.scheme()) {
       return None;
     }
-    self.resolver.resolve_redirects(&specifier)
+    self.resolver.resolve_redirects(&specifier, file_referrer)
   }
 
   /// Return `true` if the specifier can be resolved to a document.
-  pub fn exists(&self, specifier: &ModuleSpecifier) -> bool {
-    let specifier = self.resolve_document_specifier(specifier);
+  pub fn exists(
+    &self,
+    specifier: &ModuleSpecifier,
+    file_referrer: Option<&ModuleSpecifier>,
+  ) -> bool {
+    let specifier = self.resolve_document_specifier(specifier, file_referrer);
     if let Some(specifier) = specifier {
       if self.open_docs.contains_key(&specifier) {
         return true;
@@ -1043,11 +1091,38 @@ impl Documents {
   }
 
   /// Return a document for the specifier.
-  pub fn get(
+  pub fn get(&self, specifier: &ModuleSpecifier) -> Option<Arc<Document>> {
+    if let Some(document) = self.open_docs.get(specifier) {
+      Some(document.clone())
+    } else {
+      let old_doc = self
+        .file_system_docs
+        .docs
+        .get(specifier)
+        .map(|d| d.value().clone());
+      if let Some(old_doc) = old_doc {
+        self.file_system_docs.get(
+          specifier,
+          &self.resolver,
+          &self.config,
+          &self.cache,
+          old_doc.file_referrer(),
+        )
+      } else {
+        None
+      }
+    }
+  }
+
+  /// Return a document for the specifier.
+  pub fn get_or_load(
     &self,
-    original_specifier: &ModuleSpecifier,
+    specifier: &ModuleSpecifier,
+    referrer: &ModuleSpecifier,
   ) -> Option<Arc<Document>> {
-    let specifier = self.resolve_document_specifier(original_specifier)?;
+    let file_referrer = self.get_file_referrer(referrer);
+    let specifier =
+      self.resolve_document_specifier(specifier, file_referrer.as_deref())?;
     if let Some(document) = self.open_docs.get(&specifier) {
       Some(document.clone())
     } else {
@@ -1056,6 +1131,7 @@ impl Documents {
         &self.resolver,
         &self.config,
         &self.cache,
+        file_referrer.as_deref(),
       )
     }
   }
@@ -1110,6 +1186,7 @@ impl Documents {
     referrer: &ModuleSpecifier,
   ) -> Vec<Option<(ModuleSpecifier, MediaType)>> {
     let document = self.get(referrer);
+    let file_referrer = document.as_ref().and_then(|d| d.file_referrer());
     let dependencies = document.as_ref().map(|d| d.dependencies());
     let mut results = Vec::new();
     for specifier in specifiers {
@@ -1124,22 +1201,36 @@ impl Documents {
         dependencies.as_ref().and_then(|d| d.get(specifier))
       {
         if let Some(specifier) = dep.maybe_type.maybe_specifier() {
-          results.push(self.resolve_dependency(specifier, referrer));
+          results.push(self.resolve_dependency(
+            specifier,
+            referrer,
+            file_referrer,
+          ));
         } else if let Some(specifier) = dep.maybe_code.maybe_specifier() {
-          results.push(self.resolve_dependency(specifier, referrer));
+          results.push(self.resolve_dependency(
+            specifier,
+            referrer,
+            file_referrer,
+          ));
         } else {
           results.push(None);
         }
-      } else if let Ok(specifier) = self.resolver.as_graph_resolver().resolve(
-        specifier,
-        &deno_graph::Range {
-          specifier: referrer.clone(),
-          start: deno_graph::Position::zeroed(),
-          end: deno_graph::Position::zeroed(),
-        },
-        ResolutionMode::Types,
-      ) {
-        results.push(self.resolve_dependency(&specifier, referrer));
+      } else if let Ok(specifier) =
+        self.resolver.as_graph_resolver(file_referrer).resolve(
+          specifier,
+          &deno_graph::Range {
+            specifier: referrer.clone(),
+            start: deno_graph::Position::zeroed(),
+            end: deno_graph::Position::zeroed(),
+          },
+          ResolutionMode::Types,
+        )
+      {
+        results.push(self.resolve_dependency(
+          &specifier,
+          referrer,
+          file_referrer,
+        ));
       } else {
         results.push(None);
       }
@@ -1203,6 +1294,7 @@ impl Documents {
             &self.resolver,
             &self.config,
             &self.cache,
+            None,
           );
         }
       }
@@ -1283,6 +1375,7 @@ impl Documents {
     &self,
     specifier: &ModuleSpecifier,
     referrer: &ModuleSpecifier,
+    file_referrer: Option<&ModuleSpecifier>,
   ) -> Option<(ModuleSpecifier, MediaType)> {
     if let Some(module_name) = specifier.as_str().strip_prefix("node:") {
       if deno_node::is_builtin_node_module(module_name) {
@@ -1294,13 +1387,15 @@ impl Documents {
     }
 
     if let Ok(npm_ref) = NpmPackageReqReference::from_specifier(specifier) {
-      return self.resolver.npm_to_file_url(&npm_ref, referrer);
+      return self
+        .resolver
+        .npm_to_file_url(&npm_ref, referrer, file_referrer);
     }
-    let Some(doc) = self.get(specifier) else {
+    let Some(doc) = self.get_or_load(specifier, referrer) else {
       return Some((specifier.clone(), MediaType::from_specifier(specifier)));
     };
-    if let Some(specifier) = doc.maybe_types_dependency().maybe_specifier() {
-      self.resolve_dependency(specifier, referrer)
+    if let Some(types) = doc.maybe_types_dependency().maybe_specifier() {
+      self.resolve_dependency(types, specifier, file_referrer)
     } else {
       let media_type = doc.media_type();
       Some((doc.specifier().clone(), media_type))
@@ -1364,11 +1459,17 @@ fn parse_and_analyze_module(
   text_info: SourceTextInfo,
   maybe_headers: Option<&HashMap<String, String>>,
   media_type: MediaType,
+  file_referrer: Option<&ModuleSpecifier>,
   resolver: &LspResolver,
 ) -> (Option<ParsedSourceResult>, Option<ModuleResult>) {
   let parsed_source_result = parse_source(specifier, text_info, media_type);
-  let module_result =
-    analyze_module(specifier, &parsed_source_result, maybe_headers, resolver);
+  let module_result = analyze_module(
+    specifier,
+    &parsed_source_result,
+    maybe_headers,
+    file_referrer,
+    resolver,
+  );
   (Some(parsed_source_result), Some(module_result))
 }
 
@@ -1391,6 +1492,7 @@ fn analyze_module(
   specifier: &ModuleSpecifier,
   parsed_source_result: &ParsedSourceResult,
   maybe_headers: Option<&HashMap<String, String>>,
+  file_referrer: Option<&ModuleSpecifier>,
   resolver: &LspResolver,
 ) -> ModuleResult {
   match parsed_source_result {
@@ -1404,8 +1506,8 @@ fn analyze_module(
         // dynamic imports like import(`./dir/${something}`) in the LSP
         file_system: &deno_graph::source::NullFileSystem,
         jsr_url_provider: &CliJsrUrlProvider,
-        maybe_resolver: Some(resolver.as_graph_resolver()),
-        maybe_npm_resolver: Some(resolver.as_graph_npm_resolver()),
+        maybe_resolver: Some(resolver.as_graph_resolver(file_referrer)),
+        maybe_npm_resolver: Some(resolver.as_graph_npm_resolver(file_referrer)),
       },
     )),
     Err(err) => Err(deno_graph::ModuleGraphError::ModuleError(
@@ -1448,6 +1550,7 @@ console.log(b);
       1,
       "javascript".parse().unwrap(),
       content.into(),
+      None,
     );
     assert!(document.is_diagnosable());
     assert!(document.is_open());
@@ -1473,6 +1576,7 @@ console.log(b);
       1,
       "javascript".parse().unwrap(),
       content.into(),
+      None,
     );
     documents
       .change(
@@ -1517,6 +1621,7 @@ console.log(b, "hello deno");
       1,
       LanguageId::TypeScript,
       "".into(),
+      None,
     );
 
     // make a clone of the document store and close the document in that one
@@ -1586,6 +1691,7 @@ console.log(b, "hello deno");
         1,
         LanguageId::TypeScript,
         "import {} from 'test';".into(),
+        None,
       );
 
       assert_eq!(
diff --git a/cli/lsp/language_server.rs b/cli/lsp/language_server.rs
index 009b46654..17ed02cd6 100644
--- a/cli/lsp/language_server.rs
+++ b/cli/lsp/language_server.rs
@@ -114,7 +114,10 @@ impl RootCertStoreProvider for LspRootCertStoreProvider {
 }
 
 #[derive(Debug, Clone)]
-pub struct LanguageServer(Arc<tokio::sync::RwLock<Inner>>, CancellationToken);
+pub struct LanguageServer(
+  pub Arc<tokio::sync::RwLock<Inner>>,
+  CancellationToken,
+);
 
 /// Snapshot of the state used by TSC.
 #[derive(Clone, Debug, Default)]
@@ -122,7 +125,7 @@ pub struct StateSnapshot {
   pub project_version: usize,
   pub assets: AssetsSnapshot,
   pub config: Arc<Config>,
-  pub documents: Documents,
+  pub documents: Arc<Documents>,
   pub resolver: Arc<LspResolver>,
 }
 
@@ -149,7 +152,7 @@ impl Default for LanguageServerTaskQueue {
 }
 
 impl LanguageServerTaskQueue {
-  fn queue_task(&self, task_fn: LanguageServerTaskFn) -> bool {
+  pub fn queue_task(&self, task_fn: LanguageServerTaskFn) -> bool {
     self.task_tx.send(task_fn).is_ok()
   }
 
@@ -578,7 +581,7 @@ impl Inner {
       project_version: self.project_version,
       assets: self.assets.snapshot(),
       config: Arc::new(self.config.clone()),
-      documents: self.documents.clone(),
+      documents: Arc::new(self.documents.clone()),
       resolver: self.resolver.snapshot(),
     })
   }
@@ -990,6 +993,8 @@ impl Inner {
         params.text_document.uri
       );
     }
+    let file_referrer = (params.text_document.uri.scheme() == "file")
+      .then(|| params.text_document.uri.clone());
     let specifier = self
       .url_map
       .normalize_url(&params.text_document.uri, LspUrlKind::File);
@@ -998,6 +1003,7 @@ impl Inner {
       params.text_document.version,
       params.text_document.language_id.parse().unwrap(),
       params.text_document.text.into(),
+      file_referrer,
     );
     self.project_changed([(document.specifier(), ChangeKind::Opened)], false);
     if document.is_diagnosable() {
@@ -1228,6 +1234,8 @@ impl Inner {
     &self,
     params: DocumentFormattingParams,
   ) -> LspResult<Option<Vec<TextEdit>>> {
+    let file_referrer = (params.text_document.uri.scheme() == "file")
+      .then(|| params.text_document.uri.clone());
     let mut specifier = self
       .url_map
       .normalize_url(&params.text_document.uri, LspUrlKind::File);
@@ -1241,7 +1249,9 @@ impl Inner {
     {
       return Ok(None);
     }
-    let document = match self.documents.get(&specifier) {
+    let document =
+      file_referrer.and_then(|r| self.documents.get_or_load(&specifier, &r));
+    let document = match document {
       Some(doc) if doc.is_open() => doc,
       _ => return Ok(None),
     };
@@ -1329,41 +1339,44 @@ impl Inner {
 
     let mark = self.performance.mark_with_args("lsp.hover", &params);
     let asset_or_doc = self.get_asset_or_document(&specifier)?;
+    let file_referrer = asset_or_doc.document().and_then(|d| d.file_referrer());
     let hover = if let Some((_, dep, range)) = asset_or_doc
       .get_maybe_dependency(&params.text_document_position_params.position)
     {
-      let dep_doc = dep.get_code().and_then(|s| self.documents.get(s));
+      let dep_doc = dep
+        .get_code()
+        .and_then(|s| self.documents.get_or_load(s, &specifier));
       let dep_maybe_types_dependency =
         dep_doc.as_ref().map(|d| d.maybe_types_dependency());
       let value = match (dep.maybe_code.is_none(), dep.maybe_type.is_none(), &dep_maybe_types_dependency) {
         (false, false, None) => format!(
           "**Resolved Dependency**\n\n**Code**: {}\n\n**Types**: {}\n",
-          self.resolution_to_hover_text(&dep.maybe_code),
-          self.resolution_to_hover_text(&dep.maybe_type),
+          self.resolution_to_hover_text(&dep.maybe_code, file_referrer),
+          self.resolution_to_hover_text(&dep.maybe_type, file_referrer),
         ),
         (false, false, Some(types_dep)) if !types_dep.is_none() => format!(
           "**Resolved Dependency**\n\n**Code**: {}\n**Types**: {}\n**Import Types**: {}\n",
-          self.resolution_to_hover_text(&dep.maybe_code),
-          self.resolution_to_hover_text(&dep.maybe_type),
-          self.resolution_to_hover_text(types_dep),
+          self.resolution_to_hover_text(&dep.maybe_code, file_referrer),
+          self.resolution_to_hover_text(&dep.maybe_type, file_referrer),
+          self.resolution_to_hover_text(types_dep, file_referrer),
         ),
         (false, false, Some(_)) => format!(
           "**Resolved Dependency**\n\n**Code**: {}\n\n**Types**: {}\n",
-          self.resolution_to_hover_text(&dep.maybe_code),
-          self.resolution_to_hover_text(&dep.maybe_type),
+          self.resolution_to_hover_text(&dep.maybe_code, file_referrer),
+          self.resolution_to_hover_text(&dep.maybe_type, file_referrer),
         ),
         (false, true, Some(types_dep)) if !types_dep.is_none() => format!(
           "**Resolved Dependency**\n\n**Code**: {}\n\n**Types**: {}\n",
-          self.resolution_to_hover_text(&dep.maybe_code),
-          self.resolution_to_hover_text(types_dep),
+          self.resolution_to_hover_text(&dep.maybe_code, file_referrer),
+          self.resolution_to_hover_text(types_dep, file_referrer),
         ),
         (false, true, _) => format!(
           "**Resolved Dependency**\n\n**Code**: {}\n",
-          self.resolution_to_hover_text(&dep.maybe_code),
+          self.resolution_to_hover_text(&dep.maybe_code, file_referrer),
         ),
         (true, false, _) => format!(
           "**Resolved Dependency**\n\n**Types**: {}\n",
-          self.resolution_to_hover_text(&dep.maybe_type),
+          self.resolution_to_hover_text(&dep.maybe_type, file_referrer),
         ),
         (true, true, _) => unreachable!("{}", json!(params)),
       };
@@ -1394,7 +1407,11 @@ impl Inner {
     Ok(hover)
   }
 
-  fn resolution_to_hover_text(&self, resolution: &Resolution) -> String {
+  fn resolution_to_hover_text(
+    &self,
+    resolution: &Resolution,
+    file_referrer: Option<&ModuleSpecifier>,
+  ) -> String {
     match resolution {
       Resolution::Ok(resolved) => {
         let specifier = &resolved.specifier;
@@ -1416,7 +1433,9 @@ impl Inner {
             if let Ok(jsr_req_ref) =
               JsrPackageReqReference::from_specifier(specifier)
             {
-              if let Some(url) = self.resolver.jsr_to_registry_url(&jsr_req_ref)
+              if let Some(url) = self
+                .resolver
+                .jsr_to_registry_url(&jsr_req_ref, file_referrer)
               {
                 result = format!("{result} (<{url}>)");
               }
diff --git a/cli/lsp/resolver.rs b/cli/lsp/resolver.rs
index 27c4f6acf..c4d97177f 100644
--- a/cli/lsp/resolver.rs
+++ b/cli/lsp/resolver.rs
@@ -184,31 +184,49 @@ impl LspResolver {
     Ok(())
   }
 
-  pub fn as_graph_resolver(&self) -> &dyn Resolver {
+  pub fn as_graph_resolver(
+    &self,
+    _file_referrer: Option<&ModuleSpecifier>,
+  ) -> &dyn Resolver {
     self.graph_resolver.as_ref()
   }
 
-  pub fn as_graph_npm_resolver(&self) -> &dyn NpmResolver {
+  pub fn as_graph_npm_resolver(
+    &self,
+    _file_referrer: Option<&ModuleSpecifier>,
+  ) -> &dyn NpmResolver {
     self.graph_resolver.as_ref()
   }
 
-  pub fn maybe_managed_npm_resolver(&self) -> Option<&ManagedCliNpmResolver> {
+  pub fn maybe_managed_npm_resolver(
+    &self,
+    _file_referrer: Option<&ModuleSpecifier>,
+  ) -> Option<&ManagedCliNpmResolver> {
     self.npm_resolver.as_ref().and_then(|r| r.as_managed())
   }
 
-  pub fn graph_import_specifiers(
+  pub fn graph_imports_by_referrer(
     &self,
-  ) -> impl Iterator<Item = &ModuleSpecifier> {
+  ) -> IndexMap<&ModuleSpecifier, Vec<&ModuleSpecifier>> {
     self
       .graph_imports
-      .values()
-      .flat_map(|i| i.dependencies.values())
-      .flat_map(|value| value.get_type().or_else(|| value.get_code()))
+      .iter()
+      .map(|(s, i)| {
+        (
+          s,
+          i.dependencies
+            .values()
+            .flat_map(|d| d.get_type().or_else(|| d.get_code()))
+            .collect(),
+        )
+      })
+      .collect()
   }
 
   pub fn jsr_to_registry_url(
     &self,
     req_ref: &JsrPackageReqReference,
+    _file_referrer: Option<&ModuleSpecifier>,
   ) -> Option<ModuleSpecifier> {
     self.jsr_resolver.as_ref()?.jsr_to_registry_url(req_ref)
   }
@@ -217,11 +235,16 @@ impl LspResolver {
     &self,
     nv: &PackageNv,
     path: &str,
+    _file_referrer: Option<&ModuleSpecifier>,
   ) -> Option<String> {
     self.jsr_resolver.as_ref()?.lookup_export_for_path(nv, path)
   }
 
-  pub fn jsr_lookup_req_for_nv(&self, nv: &PackageNv) -> Option<PackageReq> {
+  pub fn jsr_lookup_req_for_nv(
+    &self,
+    nv: &PackageNv,
+    _file_referrer: Option<&ModuleSpecifier>,
+  ) -> Option<PackageReq> {
     self.jsr_resolver.as_ref()?.lookup_req_for_nv(nv)
   }
 
@@ -229,6 +252,7 @@ impl LspResolver {
     &self,
     req_ref: &NpmPackageReqReference,
     referrer: &ModuleSpecifier,
+    _file_referrer: Option<&ModuleSpecifier>,
   ) -> Option<(ModuleSpecifier, MediaType)> {
     let node_resolver = self.node_resolver.as_ref()?;
     Some(NodeResolution::into_specifier_and_media_type(
@@ -275,6 +299,7 @@ impl LspResolver {
   pub fn resolve_redirects(
     &self,
     specifier: &ModuleSpecifier,
+    _file_referrer: Option<&ModuleSpecifier>,
   ) -> Option<ModuleSpecifier> {
     let Some(redirect_resolver) = self.redirect_resolver.as_ref() else {
       return Some(specifier.clone());
@@ -285,6 +310,7 @@ impl LspResolver {
   pub fn redirect_chain_headers(
     &self,
     specifier: &ModuleSpecifier,
+    _file_referrer: Option<&ModuleSpecifier>,
   ) -> Vec<(ModuleSpecifier, Arc<HashMap<String, String>>)> {
     let Some(redirect_resolver) = self.redirect_resolver.as_ref() else {
       return vec![];
diff --git a/cli/lsp/tsc.rs b/cli/lsp/tsc.rs
index 5d5f5228c..9215be074 100644
--- a/cli/lsp/tsc.rs
+++ b/cli/lsp/tsc.rs
@@ -4,6 +4,7 @@ use super::analysis::CodeActionData;
 use super::code_lens;
 use super::config;
 use super::documents::AssetOrDocument;
+use super::documents::Document;
 use super::documents::DocumentsFilter;
 use super::language_server;
 use super::language_server::StateSnapshot;
@@ -393,8 +394,8 @@ impl TsServer {
     let _join_handle = thread::spawn(move || {
       run_tsc_thread(
         receiver,
-        performance.clone(),
-        specifier_map.clone(),
+        performance,
+        specifier_map,
         maybe_inspector_server,
       )
     });
@@ -3990,6 +3991,8 @@ struct State {
   response_tx: Option<oneshot::Sender<Result<String, AnyError>>>,
   state_snapshot: Arc<StateSnapshot>,
   specifier_map: Arc<TscSpecifierMap>,
+  root_referrers: HashMap<ModuleSpecifier, ModuleSpecifier>,
+  last_referrer: Option<ModuleSpecifier>,
   token: CancellationToken,
   pending_requests: Option<UnboundedReceiver<Request>>,
   mark: Option<PerformanceMark>,
@@ -4008,12 +4011,30 @@ impl State {
       response_tx: None,
       state_snapshot,
       specifier_map,
+      root_referrers: Default::default(),
+      last_referrer: None,
       token: Default::default(),
       mark: None,
       pending_requests: Some(pending_requests),
     }
   }
 
+  fn get_document(&self, specifier: &ModuleSpecifier) -> Option<Arc<Document>> {
+    if let Some(referrer) = self.root_referrers.get(specifier) {
+      self
+        .state_snapshot
+        .documents
+        .get_or_load(specifier, referrer)
+    } else if let Some(referrer) = &self.last_referrer {
+      self
+        .state_snapshot
+        .documents
+        .get_or_load(specifier, referrer)
+    } else {
+      self.state_snapshot.documents.get(specifier)
+    }
+  }
+
   fn get_asset_or_document(
     &self,
     specifier: &ModuleSpecifier,
@@ -4022,10 +4043,8 @@ impl State {
     if specifier.scheme() == "asset" {
       snapshot.assets.get(specifier).map(AssetOrDocument::Asset)
     } else {
-      snapshot
-        .documents
-        .get(specifier)
-        .map(AssetOrDocument::Document)
+      let document = self.get_document(specifier);
+      document.map(AssetOrDocument::Document)
     }
   }
 
@@ -4037,11 +4056,8 @@ impl State {
         None
       }
     } else {
-      self
-        .state_snapshot
-        .documents
-        .get(specifier)
-        .map(|d| d.script_version())
+      let document = self.get_document(specifier);
+      document.map(|d| d.script_version())
     }
   }
 }
@@ -4222,7 +4238,7 @@ fn op_resolve_inner(
       })
     })
     .collect();
-
+  state.last_referrer = Some(referrer);
   state.performance.measure(mark);
   Ok(specifiers)
 }
@@ -4235,6 +4251,7 @@ fn op_respond(
 ) {
   let state = state.borrow_mut::<State>();
   state.performance.measure(state.mark.take().unwrap());
+  state.last_referrer = None;
   let response = if !error.is_empty() {
     Err(anyhow!("tsc error: {error}"))
   } else {
@@ -4269,9 +4286,16 @@ fn op_script_names(state: &mut OpState) -> Vec<String> {
   }
 
   // inject these next because they're global
-  for specifier in state.state_snapshot.resolver.graph_import_specifiers() {
-    if seen.insert(Cow::Borrowed(specifier.as_str())) {
-      result.push(specifier.to_string());
+  for (referrer, specifiers) in
+    state.state_snapshot.resolver.graph_imports_by_referrer()
+  {
+    for specifier in specifiers {
+      if seen.insert(Cow::Borrowed(specifier.as_str())) {
+        result.push(specifier.to_string());
+      }
+      state
+        .root_referrers
+        .insert(specifier.clone(), referrer.clone());
     }
   }
 
@@ -4289,8 +4313,12 @@ fn op_script_names(state: &mut OpState) -> Vec<String> {
       let types_specifier = (|| {
         let documents = &state.state_snapshot.documents;
         let types = doc.maybe_types_dependency().maybe_specifier()?;
-        let (types, _) = documents.resolve_dependency(types, specifier)?;
-        let types_doc = documents.get(&types)?;
+        let (types, _) = documents.resolve_dependency(
+          types,
+          specifier,
+          doc.file_referrer(),
+        )?;
+        let types_doc = documents.get_or_load(&types, specifier)?;
         Some(types_doc.specifier().clone())
       })();
       // If there is a types dep, use that as the root instead. But if the doc
@@ -5169,11 +5197,12 @@ mod tests {
         *version,
         *language_id,
         (*source).into(),
+        None,
       );
     }
     let snapshot = Arc::new(StateSnapshot {
       project_version: 0,
-      documents,
+      documents: Arc::new(documents),
       assets: Default::default(),
       config: Arc::new(config),
       resolver,