perf(lsp): only store parsed sources for open documents (#23454)

8 files changed, 199 insertions, 118 deletions

diff --git a/cli/lsp/config.rs b/cli/lsp/config.rs
index ec736f365..c58d2e86d 100644
--- a/cli/lsp/config.rs
+++ b/cli/lsp/config.rs
@@ -870,7 +870,7 @@ impl Settings {
   }
 }
 
-#[derive(Debug, Default)]
+#[derive(Clone, Debug, Default)]
 pub struct Config {
   pub client_capabilities: ClientCapabilities,
   pub settings: Settings,
diff --git a/cli/lsp/documents.rs b/cli/lsp/documents.rs
index 0bcd9a8c8..207a7241b 100644
--- a/cli/lsp/documents.rs
+++ b/cli/lsp/documents.rs
@@ -5,6 +5,8 @@ use super::cache::calculate_fs_version_at_path;
 use super::cache::LSP_DISALLOW_GLOBAL_TO_LOCAL_COPY;
 use super::config::Config;
 use super::language_server::StateNpmSnapshot;
+use super::testing::TestCollector;
+use super::testing::TestModule;
 use super::text::LineIndex;
 use super::tsc;
 use super::tsc::AssetDocument;
@@ -23,12 +25,14 @@ use crate::resolver::SloppyImportsResolver;
 use deno_runtime::fs_util::specifier_to_file_path;
 
 use dashmap::DashMap;
+use deno_ast::swc::visit::VisitWith;
 use deno_ast::MediaType;
 use deno_ast::ParsedSource;
 use deno_ast::SourceTextInfo;
 use deno_core::error::custom_error;
 use deno_core::error::AnyError;
 use deno_core::futures::future;
+use deno_core::futures::future::Shared;
 use deno_core::futures::FutureExt;
 use deno_core::parking_lot::Mutex;
 use deno_core::ModuleSpecifier;
@@ -50,7 +54,9 @@ use std::collections::BTreeSet;
 use std::collections::HashMap;
 use std::collections::HashSet;
 use std::fs;
+use std::future::Future;
 use std::ops::Range;
+use std::pin::Pin;
 use std::str::FromStr;
 use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
@@ -213,6 +219,62 @@ impl AssetOrDocument {
 
 type ModuleResult = Result<deno_graph::JsModule, deno_graph::ModuleGraphError>;
 type ParsedSourceResult = Result<ParsedSource, deno_ast::ParseDiagnostic>;
+type TestModuleFut =
+  Shared<Pin<Box<dyn Future<Output = Option<Arc<TestModule>>> + Send>>>;
+
+fn media_type_is_diagnosable(media_type: MediaType) -> bool {
+  matches!(
+    media_type,
+    MediaType::JavaScript
+      | MediaType::Jsx
+      | MediaType::Mjs
+      | MediaType::Cjs
+      | MediaType::TypeScript
+      | MediaType::Tsx
+      | MediaType::Mts
+      | MediaType::Cts
+      | MediaType::Dts
+      | MediaType::Dmts
+      | MediaType::Dcts
+  )
+}
+
+fn get_maybe_test_module_fut(
+  maybe_parsed_source: Option<&ParsedSourceResult>,
+  config: &Config,
+) -> Option<TestModuleFut> {
+  if !config.client_capabilities.testing_api {
+    return None;
+  }
+  let parsed_source = maybe_parsed_source?.as_ref().ok()?.clone();
+  let specifier = parsed_source.specifier();
+  if specifier.scheme() != "file" {
+    return None;
+  }
+  if !media_type_is_diagnosable(parsed_source.media_type()) {
+    return None;
+  }
+  if !config.specifier_enabled_for_test(specifier) {
+    return None;
+  }
+  let handle = tokio::task::spawn_blocking(move || {
+    let mut collector = TestCollector::new(
+      parsed_source.specifier().clone(),
+      parsed_source.text_info().clone(),
+    );
+    parsed_source.module().visit_with(&mut collector);
+    Arc::new(collector.take())
+  })
+  .map(Result::ok)
+  .boxed()
+  .shared();
+  Some(handle)
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct DocumentOpenData {
+  maybe_parsed_source: Option<ParsedSourceResult>,
+}
 
 #[derive(Debug)]
 pub struct Document {
@@ -227,13 +289,16 @@ pub struct Document {
   // this is a lazily constructed value based on the state of the document,
   // so having a mutex to hold it is ok
   maybe_navigation_tree: Mutex<Option<Arc<tsc::NavigationTree>>>,
-  maybe_parsed_source: Option<ParsedSourceResult>,
+  maybe_test_module_fut: Option<TestModuleFut>,
   media_type: MediaType,
+  /// Present if and only if this is an open document.
+  open_data: Option<DocumentOpenData>,
   specifier: ModuleSpecifier,
   text_info: SourceTextInfo,
 }
 
 impl Document {
+  #[allow(clippy::too_many_arguments)]
   fn new(
     specifier: ModuleSpecifier,
     fs_version: String,
@@ -242,6 +307,7 @@ impl Document {
     resolver: &dyn deno_graph::source::Resolver,
     maybe_node_resolver: Option<&CliNodeResolver>,
     npm_resolver: &dyn deno_graph::source::NpmResolver,
+    config: &Config,
   ) -> Arc<Self> {
     // we only ever do `Document::new` on disk resources that are supposed to
     // be diagnosable, unlike `Document::open`, so it is safe to unconditionally
@@ -269,6 +335,8 @@ impl Document {
       .as_ref()
       .and_then(|m| Some(Arc::new(m.maybe_types_dependency.clone()?)));
     let line_index = Arc::new(LineIndex::new(text_info.text_str()));
+    let maybe_test_module_fut =
+      get_maybe_test_module_fut(maybe_parsed_source.as_ref(), config);
     Arc::new(Document {
       dependencies,
       maybe_types_dependency,
@@ -278,9 +346,9 @@ impl Document {
       maybe_language_id: None,
       maybe_lsp_version: None,
       maybe_navigation_tree: Mutex::new(None),
-      maybe_parsed_source: maybe_parsed_source
-        .filter(|_| specifier.scheme() == "file"),
+      maybe_test_module_fut,
       media_type,
+      open_data: None,
       text_info,
       specifier,
     })
@@ -291,6 +359,7 @@ impl Document {
     resolver: &dyn deno_graph::source::Resolver,
     maybe_node_resolver: Option<&CliNodeResolver>,
     npm_resolver: &dyn deno_graph::source::NpmResolver,
+    config: &Config,
   ) -> Option<Arc<Self>> {
     let media_type = resolve_media_type(
       &self.specifier,
@@ -301,6 +370,7 @@ impl Document {
     let dependencies;
     let maybe_types_dependency;
     let maybe_parsed_source;
+    let maybe_test_module_fut;
     if media_type != self.media_type {
       let parsed_source_result =
         parse_source(&self.specifier, self.text_info.clone(), media_type);
@@ -320,6 +390,8 @@ impl Document {
         .as_ref()
         .and_then(|m| Some(Arc::new(m.maybe_types_dependency.clone()?)));
       maybe_parsed_source = Some(parsed_source_result);
+      maybe_test_module_fut =
+        get_maybe_test_module_fut(maybe_parsed_source.as_ref(), config);
     } else {
       dependencies = Arc::new(
         self
@@ -336,22 +408,28 @@ impl Document {
       maybe_types_dependency = self.maybe_types_dependency.as_ref().map(|d| {
         Arc::new(d.with_new_resolver(Some(resolver), Some(npm_resolver)))
       });
-      maybe_parsed_source = self.maybe_parsed_source.clone();
+      maybe_parsed_source = self.maybe_parsed_source();
+      maybe_test_module_fut = self
+        .maybe_test_module_fut
+        .clone()
+        .filter(|_| config.specifier_enabled_for_test(&self.specifier));
     }
     Some(Arc::new(Self {
       // updated properties
       dependencies,
       maybe_types_dependency,
       maybe_navigation_tree: Mutex::new(None),
-      maybe_parsed_source: maybe_parsed_source
-        .filter(|_| self.specifier.scheme() == "file"),
       // maintain - this should all be copies/clones
       fs_version: self.fs_version.clone(),
       line_index: self.line_index.clone(),
       maybe_headers: self.maybe_headers.clone(),
       maybe_language_id: self.maybe_language_id,
       maybe_lsp_version: self.maybe_lsp_version,
+      maybe_test_module_fut,
       media_type,
+      open_data: self.open_data.is_some().then_some(DocumentOpenData {
+        maybe_parsed_source,
+      }),
       text_info: self.text_info.clone(),
       specifier: self.specifier.clone(),
     }))
@@ -368,6 +446,7 @@ impl Document {
     resolver: &dyn deno_graph::source::Resolver,
     maybe_node_resolver: Option<&CliNodeResolver>,
     npm_resolver: &dyn deno_graph::source::NpmResolver,
+    config: &Config,
   ) -> Arc<Self> {
     let text_info = SourceTextInfo::new(content);
     let media_type = resolve_media_type(
@@ -397,6 +476,8 @@ impl Document {
       .as_ref()
       .and_then(|m| Some(Arc::new(m.maybe_types_dependency.clone()?)));
     let line_index = Arc::new(LineIndex::new(text_info.text_str()));
+    let maybe_test_module_fut =
+      get_maybe_test_module_fut(maybe_parsed_source.as_ref(), config);
     Arc::new(Self {
       dependencies,
       maybe_types_dependency,
@@ -407,9 +488,11 @@ impl Document {
       maybe_lsp_version: Some(version),
       maybe_headers,
       maybe_navigation_tree: Mutex::new(None),
-      maybe_parsed_source: maybe_parsed_source
-        .filter(|_| specifier.scheme() == "file"),
+      maybe_test_module_fut,
       media_type,
+      open_data: Some(DocumentOpenData {
+        maybe_parsed_source,
+      }),
       text_info,
       specifier,
     })
@@ -421,6 +504,7 @@ impl Document {
     changes: Vec<lsp::TextDocumentContentChangeEvent>,
     resolver: &dyn deno_graph::source::Resolver,
     npm_resolver: &dyn deno_graph::source::NpmResolver,
+    config: &Config,
   ) -> Result<Arc<Self>, AnyError> {
     let mut content = self.text_info.text_str().to_string();
     let mut line_index = self.line_index.clone();
@@ -471,6 +555,8 @@ impl Document {
     } else {
       Arc::new(LineIndex::new(text_info.text_str()))
     };
+    let maybe_test_module_fut =
+      get_maybe_test_module_fut(maybe_parsed_source.as_ref(), config);
     Ok(Arc::new(Self {
       specifier: self.specifier.clone(),
       fs_version: self.fs_version.clone(),
@@ -480,11 +566,13 @@ impl Document {
       text_info,
       line_index,
       maybe_headers: self.maybe_headers.clone(),
-      maybe_parsed_source: maybe_parsed_source
-        .filter(|_| self.specifier.scheme() == "file"),
       maybe_lsp_version: Some(version),
       maybe_navigation_tree: Mutex::new(None),
+      maybe_test_module_fut,
       media_type,
+      open_data: self.open_data.is_some().then_some(DocumentOpenData {
+        maybe_parsed_source,
+      }),
     }))
   }
 
@@ -499,10 +587,11 @@ impl Document {
       text_info: self.text_info.clone(),
       line_index: self.line_index.clone(),
       maybe_headers: self.maybe_headers.clone(),
-      maybe_parsed_source: self.maybe_parsed_source.clone(),
       maybe_lsp_version: self.maybe_lsp_version,
       maybe_navigation_tree: Mutex::new(None),
+      maybe_test_module_fut: self.maybe_test_module_fut.clone(),
       media_type: self.media_type,
+      open_data: self.open_data.clone(),
     })
   }
 
@@ -534,20 +623,7 @@ impl Document {
   }
 
   pub fn is_diagnosable(&self) -> bool {
-    matches!(
-      self.media_type(),
-      MediaType::JavaScript
-        | MediaType::Jsx
-        | MediaType::Mjs
-        | MediaType::Cjs
-        | MediaType::TypeScript
-        | MediaType::Tsx
-        | MediaType::Mts
-        | MediaType::Cts
-        | MediaType::Dts
-        | MediaType::Dmts
-        | MediaType::Dcts
-    )
+    media_type_is_diagnosable(self.media_type())
   }
 
   pub fn is_open(&self) -> bool {
@@ -578,7 +654,11 @@ impl Document {
   pub fn maybe_parsed_source(
     &self,
   ) -> Option<Result<deno_ast::ParsedSource, deno_ast::ParseDiagnostic>> {
-    self.maybe_parsed_source.clone()
+    self.open_data.as_ref()?.maybe_parsed_source.clone()
+  }
+
+  pub async fn maybe_test_module(&self) -> Option<Arc<TestModule>> {
+    self.maybe_test_module_fut.clone()?.await
   }
 
   pub fn maybe_navigation_tree(&self) -> Option<Arc<tsc::NavigationTree>> {
@@ -740,6 +820,7 @@ impl FileSystemDocuments {
     specifier: &ModuleSpecifier,
     maybe_node_resolver: Option<&CliNodeResolver>,
     npm_resolver: &dyn deno_graph::source::NpmResolver,
+    config: &Config,
   ) -> Option<Arc<Document>> {
     let fs_version = if specifier.scheme() == "data" {
       Some("1".to_string())
@@ -757,6 +838,7 @@ impl FileSystemDocuments {
         specifier,
         maybe_node_resolver,
         npm_resolver,
+        config,
       )
     } else {
       file_system_doc
@@ -772,6 +854,7 @@ impl FileSystemDocuments {
     specifier: &ModuleSpecifier,
     maybe_node_resolver: Option<&CliNodeResolver>,
     npm_resolver: &dyn deno_graph::source::NpmResolver,
+    config: &Config,
   ) -> Option<Arc<Document>> {
     let doc = if specifier.scheme() == "file" {
       let path = specifier_to_file_path(specifier).ok()?;
@@ -787,6 +870,7 @@ impl FileSystemDocuments {
         resolver,
         maybe_node_resolver,
         npm_resolver,
+        config,
       )
     } else if specifier.scheme() == "data" {
       let source = deno_graph::source::RawDataUrl::parse(specifier)
@@ -801,6 +885,7 @@ impl FileSystemDocuments {
         resolver,
         maybe_node_resolver,
         npm_resolver,
+        config,
       )
     } else {
       let fs_version = calculate_fs_version(cache, specifier)?;
@@ -829,6 +914,7 @@ impl FileSystemDocuments {
         resolver,
         maybe_node_resolver,
         npm_resolver,
+        config,
       )
     };
     self.docs.insert(specifier.clone(), doc.clone());
@@ -864,6 +950,7 @@ pub enum DocumentsFilter {
 pub struct Documents {
   /// The DENO_DIR that the documents looks for non-file based modules.
   cache: Arc<dyn HttpCache>,
+  config: Arc<Config>,
   /// A flag that indicates that stated data is potentially invalid and needs to
   /// be recalculated before being considered valid.
   dirty: bool,
@@ -896,6 +983,7 @@ impl Documents {
   pub fn new(cache: Arc<dyn HttpCache>) -> Self {
     Self {
       cache: cache.clone(),
+      config: Default::default(),
       dirty: true,
       open_docs: HashMap::default(),
       file_system_docs: Default::default(),
@@ -920,6 +1008,10 @@ impl Documents {
     }
   }
 
+  pub fn initialize(&mut self, config: &Config) {
+    self.config = Arc::new(config.clone());
+  }
+
   pub fn module_graph_imports(&self) -> impl Iterator<Item = &ModuleSpecifier> {
     self
       .imports
@@ -954,6 +1046,7 @@ impl Documents {
       resolver,
       self.maybe_node_resolver.as_deref(),
       npm_resolver,
+      self.config.as_ref(),
     );
 
     self.file_system_docs.remove_document(&specifier);
@@ -989,6 +1082,7 @@ impl Documents {
       changes,
       self.get_resolver(),
       self.get_npm_resolver(),
+      self.config.as_ref(),
     )?;
     self.open_docs.insert(doc.specifier().clone(), doc.clone());
     Ok(doc)
@@ -1154,6 +1248,7 @@ impl Documents {
         &specifier,
         self.maybe_node_resolver.as_deref(),
         self.get_npm_resolver(),
+        self.config.as_ref(),
       )
     }
   }
@@ -1331,6 +1426,7 @@ impl Documents {
     npm_resolver: Option<Arc<dyn CliNpmResolver>>,
     workspace_files: &BTreeSet<ModuleSpecifier>,
   ) {
+    self.config = Arc::new(config.clone());
     let config_data = config.tree.root_data();
     let config_file = config_data.and_then(|d| d.config_file.as_deref());
     self.maybe_node_resolver = node_resolver.clone();
@@ -1409,6 +1505,7 @@ impl Documents {
           resolver,
           self.maybe_node_resolver.as_deref(),
           npm_resolver,
+          self.config.as_ref(),
         ) {
           *doc = new_doc;
         }
@@ -1421,6 +1518,7 @@ impl Documents {
           resolver,
           self.maybe_node_resolver.as_deref(),
           npm_resolver,
+          self.config.as_ref(),
         ) {
           *doc.value_mut() = new_doc;
         }
@@ -1444,6 +1542,7 @@ impl Documents {
             specifier,
             self.maybe_node_resolver.as_deref(),
             npm_resolver,
+            self.config.as_ref(),
           );
         }
       }
diff --git a/cli/lsp/language_server.rs b/cli/lsp/language_server.rs
index d2e56774f..2529df6f7 100644
--- a/cli/lsp/language_server.rs
+++ b/cli/lsp/language_server.rs
@@ -404,8 +404,11 @@ impl LanguageServer {
     let inner = self.0.read().await;
     if let Some(testing_server) = &inner.maybe_testing_server {
       match params.map(serde_json::from_value) {
-        Some(Ok(params)) => testing_server
-          .run_request(params, inner.config.workspace_settings().clone()),
+        Some(Ok(params)) => {
+          testing_server
+            .run_request(params, inner.config.workspace_settings().clone())
+            .await
+        }
         Some(Err(err)) => Err(LspError::invalid_params(err.to_string())),
         None => Err(LspError::invalid_params("Missing parameters")),
       }
@@ -972,6 +975,8 @@ impl Inner {
       self.config.update_capabilities(&params.capabilities);
     }
 
+    self.documents.initialize(&self.config);
+
     if let Err(e) = self
       .ts_server
       .start(self.config.internal_inspect().to_address())
diff --git a/cli/lsp/testing/collectors.rs b/cli/lsp/testing/collectors.rs
index 8579ccc7d..508e50f9b 100644
--- a/cli/lsp/testing/collectors.rs
+++ b/cli/lsp/testing/collectors.rs
@@ -451,13 +451,9 @@ pub struct TestCollector {
 }
 
 impl TestCollector {
-  pub fn new(
-    specifier: ModuleSpecifier,
-    script_version: String,
-    text_info: SourceTextInfo,
-  ) -> Self {
+  pub fn new(specifier: ModuleSpecifier, text_info: SourceTextInfo) -> Self {
     Self {
-      test_module: TestModule::new(specifier, script_version),
+      test_module: TestModule::new(specifier),
       vars: HashSet::new(),
       fns: HashMap::new(),
       text_info,
@@ -653,8 +649,7 @@ pub mod tests {
     })
     .unwrap();
     let text_info = parsed_module.text_info().clone();
-    let mut collector =
-      TestCollector::new(specifier, "1".to_string(), text_info);
+    let mut collector = TestCollector::new(specifier, text_info);
     parsed_module.module().visit_with(&mut collector);
     collector.take()
   }
@@ -671,7 +666,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![(
           "4ebb361c93f76a0f1bac300638675609f1cf481e6f3b9006c3c98604b3a184e9"
             .to_string(),
@@ -704,7 +698,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![(
           "4ebb361c93f76a0f1bac300638675609f1cf481e6f3b9006c3c98604b3a184e9"
             .to_string(),
@@ -747,7 +740,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![
           (
             "4ebb361c93f76a0f1bac300638675609f1cf481e6f3b9006c3c98604b3a184e9".to_string(),
@@ -809,7 +801,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![
           (
             "4ebb361c93f76a0f1bac300638675609f1cf481e6f3b9006c3c98604b3a184e9".to_string(),
@@ -862,7 +853,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![(
           "4ebb361c93f76a0f1bac300638675609f1cf481e6f3b9006c3c98604b3a184e9"
             .to_string(),
@@ -897,7 +887,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![
           (
             "86b4c821900e38fc89f24bceb0e45193608ab3f9d2a6019c7b6a5aceff5d7df2".to_string(),
@@ -939,7 +928,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![(
           "4ebb361c93f76a0f1bac300638675609f1cf481e6f3b9006c3c98604b3a184e9"
             .to_string(),
@@ -973,7 +961,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![(
           "4ebb361c93f76a0f1bac300638675609f1cf481e6f3b9006c3c98604b3a184e9"
             .to_string(),
@@ -1008,7 +995,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![
           (
             "87f28e06f5ddadd90a74a93b84df2e31b9edced8301b0ad4c8fbab8d806ec99d".to_string(),
@@ -1022,7 +1008,7 @@ pub mod tests {
             },
           ),
           (
-            "e0f6a73647b763f82176c98a019e54200b799a32007f9859fb782aaa9e308568".to_string(), 
+            "e0f6a73647b763f82176c98a019e54200b799a32007f9859fb782aaa9e308568".to_string(),
             TestDefinition {
               id: "e0f6a73647b763f82176c98a019e54200b799a32007f9859fb782aaa9e308568".to_string(),
               name: "someFunction".to_string(),
@@ -1063,7 +1049,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![(
           "e0f6a73647b763f82176c98a019e54200b799a32007f9859fb782aaa9e308568"
             .to_string(),
@@ -1097,7 +1082,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![(
           "6d05d6dc35548b86a1e70acaf24a5bc2dd35db686b35b685ad5931d201b4a918"
             .to_string(),
@@ -1138,7 +1122,6 @@ pub mod tests {
       &test_module,
       &TestModule {
         specifier: test_module.specifier.clone(),
-        script_version: test_module.script_version.clone(),
         defs: vec![
           (
             "3799fc549a32532145ffc8532b0cd943e025bbc19a02e2cde9be94f87bceb829".to_string(),
diff --git a/cli/lsp/testing/definitions.rs b/cli/lsp/testing/definitions.rs
index ab47beec9..43a07c2e3 100644
--- a/cli/lsp/testing/definitions.rs
+++ b/cli/lsp/testing/definitions.rs
@@ -28,16 +28,13 @@ pub struct TestDefinition {
 #[derive(Debug, Clone, PartialEq)]
 pub struct TestModule {
   pub specifier: ModuleSpecifier,
-  /// The version of the document that the discovered tests relate to.
-  pub script_version: String,
   pub defs: HashMap<String, TestDefinition>,
 }
 
 impl TestModule {
-  pub fn new(specifier: ModuleSpecifier, script_version: String) -> Self {
+  pub fn new(specifier: ModuleSpecifier) -> Self {
     Self {
       specifier,
-      script_version,
       defs: Default::default(),
     }
   }
diff --git a/cli/lsp/testing/execution.rs b/cli/lsp/testing/execution.rs
index e1189ec2c..73916d0c2 100644
--- a/cli/lsp/testing/execution.rs
+++ b/cli/lsp/testing/execution.rs
@@ -3,6 +3,7 @@
 use super::definitions::TestDefinition;
 use super::definitions::TestModule;
 use super::lsp_custom;
+use super::server::TestServerTests;
 
 use crate::args::flags_from_vec;
 use crate::args::DenoSubcommand;
@@ -21,7 +22,6 @@ use deno_core::error::JsError;
 use deno_core::futures::future;
 use deno_core::futures::stream;
 use deno_core::futures::StreamExt;
-use deno_core::parking_lot::Mutex;
 use deno_core::parking_lot::RwLock;
 use deno_core::unsync::spawn;
 use deno_core::unsync::spawn_blocking;
@@ -42,7 +42,7 @@ use tower_lsp::lsp_types as lsp;
 /// any filters to be applied to those tests
 fn as_queue_and_filters(
   params: &lsp_custom::TestRunRequestParams,
-  tests: &HashMap<ModuleSpecifier, TestModule>,
+  tests: &HashMap<ModuleSpecifier, (TestModule, String)>,
 ) -> (
   HashSet<ModuleSpecifier>,
   HashMap<ModuleSpecifier, LspTestFilter>,
@@ -52,7 +52,7 @@ fn as_queue_and_filters(
 
   if let Some(include) = &params.include {
     for item in include {
-      if let Some(test_definitions) = tests.get(&item.text_document.uri) {
+      if let Some((test_definitions, _)) = tests.get(&item.text_document.uri) {
         queue.insert(item.text_document.uri.clone());
         if let Some(id) = &item.id {
           if let Some(test) = test_definitions.get(id) {
@@ -74,7 +74,7 @@ fn as_queue_and_filters(
   }
 
   for item in &params.exclude {
-    if let Some(test_definitions) = tests.get(&item.text_document.uri) {
+    if let Some((test_definitions, _)) = tests.get(&item.text_document.uri) {
       if let Some(id) = &item.id {
         // there is no way to exclude a test step
         if item.step_id.is_none() {
@@ -91,7 +91,7 @@ fn as_queue_and_filters(
     }
   }
 
-  queue.retain(|s| !tests.get(s).unwrap().is_empty());
+  queue.retain(|s| !tests.get(s).unwrap().0.is_empty());
 
   (queue, filters)
 }
@@ -147,19 +147,19 @@ pub struct TestRun {
   kind: lsp_custom::TestRunKind,
   filters: HashMap<ModuleSpecifier, LspTestFilter>,
   queue: HashSet<ModuleSpecifier>,
-  tests: Arc<Mutex<HashMap<ModuleSpecifier, TestModule>>>,
+  tests: TestServerTests,
   token: CancellationToken,
   workspace_settings: config::WorkspaceSettings,
 }
 
 impl TestRun {
-  pub fn new(
+  pub async fn init(
     params: &lsp_custom::TestRunRequestParams,
-    tests: Arc<Mutex<HashMap<ModuleSpecifier, TestModule>>>,
+    tests: TestServerTests,
     workspace_settings: config::WorkspaceSettings,
   ) -> Self {
     let (queue, filters) = {
-      let tests = tests.lock();
+      let tests = tests.lock().await;
       as_queue_and_filters(params, &tests)
     };
 
@@ -176,13 +176,13 @@ impl TestRun {
 
   /// Provide the tests of a test run as an enqueued module which can be sent
   /// to the client to indicate tests are enqueued for testing.
-  pub fn as_enqueued(&self) -> Vec<lsp_custom::EnqueuedTestModule> {
-    let tests = self.tests.lock();
+  pub async fn as_enqueued(&self) -> Vec<lsp_custom::EnqueuedTestModule> {
+    let tests = self.tests.lock().await;
     self
       .queue
       .iter()
       .map(|s| {
-        let ids = if let Some(test_module) = tests.get(s) {
+        let ids = if let Some((test_module, _)) = tests.get(s) {
           if let Some(filter) = self.filters.get(s) {
             filter.as_ids(test_module)
           } else {
@@ -332,7 +332,7 @@ impl TestRun {
           match event {
             test::TestEvent::Register(description) => {
               for (_, description) in description.into_iter() {
-                reporter.report_register(description);
+                reporter.report_register(description).await;
                 // TODO(mmastrac): we shouldn't need to clone here - we can re-use the descriptions
                 tests.write().insert(description.id, description.clone());
               }
@@ -378,7 +378,7 @@ impl TestRun {
               summary.uncaught_errors.push((origin, error));
             }
             test::TestEvent::StepRegister(description) => {
-              reporter.report_step_register(&description);
+              reporter.report_step_register(&description).await;
               test_steps.insert(description.id, description);
             }
             test::TestEvent::StepWait(id) => {
@@ -541,7 +541,7 @@ struct LspTestReporter {
   client: Client,
   id: u32,
   maybe_root_uri: Option<ModuleSpecifier>,
-  files: Arc<Mutex<HashMap<ModuleSpecifier, TestModule>>>,
+  files: TestServerTests,
   tests: IndexMap<usize, LspTestDescription>,
   current_test: Option<usize>,
 }
@@ -551,7 +551,7 @@ impl LspTestReporter {
     run: &TestRun,
     client: Client,
     maybe_root_uri: Option<&ModuleSpecifier>,
-    files: Arc<Mutex<HashMap<ModuleSpecifier, TestModule>>>,
+    files: TestServerTests,
   ) -> Self {
     Self {
       client,
@@ -576,12 +576,12 @@ impl LspTestReporter {
 
   fn report_plan(&mut self, _plan: &test::TestPlan) {}
 
-  fn report_register(&mut self, desc: &test::TestDescription) {
-    let mut files = self.files.lock();
+  async fn report_register(&mut self, desc: &test::TestDescription) {
+    let mut files = self.files.lock().await;
     let specifier = ModuleSpecifier::parse(&desc.location.file_name).unwrap();
-    let test_module = files
+    let (test_module, _) = files
       .entry(specifier.clone())
-      .or_insert_with(|| TestModule::new(specifier, "1".to_string()));
+      .or_insert_with(|| (TestModule::new(specifier), "1".to_string()));
     let (static_id, is_new) = test_module.register_dynamic(desc);
     self.tests.insert(
       desc.id,
@@ -681,12 +681,12 @@ impl LspTestReporter {
     }
   }
 
-  fn report_step_register(&mut self, desc: &test::TestStepDescription) {
-    let mut files = self.files.lock();
+  async fn report_step_register(&mut self, desc: &test::TestStepDescription) {
+    let mut files = self.files.lock().await;
     let specifier = ModuleSpecifier::parse(&desc.location.file_name).unwrap();
-    let test_module = files
+    let (test_module, _) = files
       .entry(specifier.clone())
-      .or_insert_with(|| TestModule::new(specifier, "1".to_string()));
+      .or_insert_with(|| (TestModule::new(specifier), "1".to_string()));
     let (static_id, is_new) = test_module.register_step_dynamic(
       desc,
       self.tests.get(&desc.parent_id).unwrap().static_id(),
@@ -828,7 +828,6 @@ mod tests {
     };
     let test_module = TestModule {
       specifier: specifier.clone(),
-      script_version: "1".to_string(),
       defs: vec![
         (test_def_a.id.clone(), test_def_a.clone()),
         (test_def_b.id.clone(), test_def_b.clone()),
@@ -836,10 +835,10 @@ mod tests {
       .into_iter()
       .collect(),
     };
-    tests.insert(specifier.clone(), test_module.clone());
+    tests.insert(specifier.clone(), (test_module.clone(), "1".to_string()));
     tests.insert(
       non_test_specifier.clone(),
-      TestModule::new(non_test_specifier, "1".to_string()),
+      (TestModule::new(non_test_specifier), "1".to_string()),
     );
     let (queue, filters) = as_queue_and_filters(&params, &tests);
     assert_eq!(json!(queue), json!([specifier]));
diff --git a/cli/lsp/testing/mod.rs b/cli/lsp/testing/mod.rs
index a11d3a8cc..d285481f0 100644
--- a/cli/lsp/testing/mod.rs
+++ b/cli/lsp/testing/mod.rs
@@ -6,6 +6,8 @@ mod execution;
 pub mod lsp_custom;
 mod server;
 
+pub use collectors::TestCollector;
+pub use definitions::TestModule;
 pub use lsp_custom::TEST_RUN_CANCEL_REQUEST;
 pub use lsp_custom::TEST_RUN_REQUEST;
 pub use server::TestServer;
diff --git a/cli/lsp/testing/server.rs b/cli/lsp/testing/server.rs
index bdf238078..ff59e1033 100644
--- a/cli/lsp/testing/server.rs
+++ b/cli/lsp/testing/server.rs
@@ -1,6 +1,5 @@
 // Copyright 2018-2024 the Deno authors. All rights reserved. MIT license.
 
-use super::collectors::TestCollector;
 use super::definitions::TestModule;
 use super::execution::TestRun;
 use super::lsp_custom;
@@ -12,7 +11,6 @@ use crate::lsp::documents::DocumentsFilter;
 use crate::lsp::language_server::StateSnapshot;
 use crate::lsp::performance::Performance;
 
-use deno_ast::swc::visit::VisitWith;
 use deno_core::error::AnyError;
 use deno_core::parking_lot::Mutex;
 use deno_core::serde_json::json;
@@ -36,6 +34,9 @@ fn as_delete_notification(uri: ModuleSpecifier) -> TestingNotification {
   )
 }
 
+pub type TestServerTests =
+  Arc<tokio::sync::Mutex<HashMap<ModuleSpecifier, (TestModule, String)>>>;
+
 /// The main structure which handles requests and sends notifications related
 /// to the Testing API.
 #[derive(Debug)]
@@ -47,7 +48,7 @@ pub struct TestServer {
   /// A map of run ids to test runs
   runs: Arc<Mutex<HashMap<u32, TestRun>>>,
   /// Tests that are discovered from a versioned document
-  tests: Arc<Mutex<HashMap<ModuleSpecifier, TestModule>>>,
+  tests: TestServerTests,
   /// A channel for requesting that changes to documents be statically analyzed
   /// for tests
   update_channel: mpsc::UnboundedSender<Arc<StateSnapshot>>,
@@ -59,8 +60,7 @@ impl TestServer {
     performance: Arc<Performance>,
     maybe_root_uri: Option<ModuleSpecifier>,
   ) -> Self {
-    let tests: Arc<Mutex<HashMap<ModuleSpecifier, TestModule>>> =
-      Arc::new(Mutex::new(HashMap::new()));
+    let tests = Default::default();
 
     let (update_channel, mut update_rx) =
       mpsc::unbounded_channel::<Arc<StateSnapshot>>();
@@ -88,7 +88,7 @@ impl TestServer {
             None => break,
             Some(snapshot) => {
               let mark = performance.mark("lsp.testing_update");
-              let mut tests = tests.lock();
+              let mut tests = tests.lock().await;
               // we create a list of test modules we currently are tracking
               // eliminating any we go over when iterating over the document
               let mut keys: HashSet<ModuleSpecifier> =
@@ -106,37 +106,33 @@ impl TestServer {
                 }
                 keys.remove(specifier);
                 let script_version = document.script_version();
-                let valid = if let Some(test) = tests.get(specifier) {
-                  test.script_version == script_version
-                } else {
-                  false
-                };
+                let valid =
+                  if let Some((_, old_script_version)) = tests.get(specifier) {
+                    old_script_version == &script_version
+                  } else {
+                    false
+                  };
                 if !valid {
-                  if let Some(Ok(parsed_source)) =
-                    document.maybe_parsed_source()
-                  {
-                    let was_empty = tests
-                      .remove(specifier)
-                      .map(|tm| tm.is_empty())
-                      .unwrap_or(true);
-                    let mut collector = TestCollector::new(
-                      specifier.clone(),
-                      script_version,
-                      parsed_source.text_info().clone(),
+                  let was_empty = tests
+                    .remove(specifier)
+                    .map(|(tm, _)| tm.is_empty())
+                    .unwrap_or(true);
+                  let test_module = document
+                    .maybe_test_module()
+                    .await
+                    .map(|tm| tm.as_ref().clone())
+                    .unwrap_or_else(|| TestModule::new(specifier.clone()));
+                  if !test_module.is_empty() {
+                    client.send_test_notification(
+                      test_module.as_replace_notification(mru.as_ref()),
                     );
-                    parsed_source.module().visit_with(&mut collector);
-                    let test_module = collector.take();
-                    if !test_module.is_empty() {
-                      client.send_test_notification(
-                        test_module.as_replace_notification(mru.as_ref()),
-                      );
-                    } else if !was_empty {
-                      client.send_test_notification(as_delete_notification(
-                        specifier.clone(),
-                      ));
-                    }
-                    tests.insert(specifier.clone(), test_module);
+                  } else if !was_empty {
+                    client.send_test_notification(as_delete_notification(
+                      specifier.clone(),
+                    ));
                   }
+                  tests
+                    .insert(specifier.clone(), (test_module, script_version));
                 }
               }
               for key in keys {
@@ -205,14 +201,14 @@ impl TestServer {
   }
 
   /// A request from the client to start a test run.
-  pub fn run_request(
+  pub async fn run_request(
     &self,
     params: lsp_custom::TestRunRequestParams,
     workspace_settings: config::WorkspaceSettings,
   ) -> LspResult<Option<Value>> {
     let test_run =
-      { TestRun::new(&params, self.tests.clone(), workspace_settings) };
-    let enqueued = test_run.as_enqueued();
+      { TestRun::init(&params, self.tests.clone(), workspace_settings).await };
+    let enqueued = test_run.as_enqueued().await;
     {
       let mut runs = self.runs.lock();
       runs.insert(params.id, test_run);