feat(lsp): cache jsxImportSource automatically (#21687)

2 files changed, 138 insertions, 7 deletions

diff --git a/cli/lsp/language_server.rs b/cli/lsp/language_server.rs
index 9594e6c7e..754ccd680 100644
--- a/cli/lsp/language_server.rs
+++ b/cli/lsp/language_server.rs
@@ -1,5 +1,6 @@
 // Copyright 2018-2024 the Deno authors. All rights reserved. MIT license.
 
+use base64::Engine;
 use deno_ast::MediaType;
 use deno_core::anyhow::anyhow;
 use deno_core::anyhow::Context;
@@ -30,6 +31,9 @@ use std::env;
 use std::fmt::Write as _;
 use std::path::PathBuf;
 use std::sync::Arc;
+use tokio::sync::mpsc::unbounded_channel;
+use tokio::sync::mpsc::UnboundedReceiver;
+use tokio::sync::mpsc::UnboundedSender;
 use tokio_util::sync::CancellationToken;
 use tower_lsp::jsonrpc::Error as LspError;
 use tower_lsp::jsonrpc::Result as LspResult;
@@ -177,6 +181,44 @@ pub struct StateSnapshot {
   pub npm: Option<StateNpmSnapshot>,
 }
 
+type LanguageServerTaskFn = Box<dyn FnOnce(LanguageServer) + Send + Sync>;
+
+/// Used to queue tasks from inside of the language server lock that must be
+/// commenced from outside of it. For example, queue a request to cache a module
+/// after having loaded a config file which references it.
+#[derive(Debug)]
+struct LanguageServerTaskQueue {
+  task_tx: UnboundedSender<LanguageServerTaskFn>,
+  /// This is moved out to its own task after initializing.
+  task_rx: Option<UnboundedReceiver<LanguageServerTaskFn>>,
+}
+
+impl Default for LanguageServerTaskQueue {
+  fn default() -> Self {
+    let (task_tx, task_rx) = unbounded_channel();
+    Self {
+      task_tx,
+      task_rx: Some(task_rx),
+    }
+  }
+}
+
+impl LanguageServerTaskQueue {
+  fn queue_task(&self, task_fn: LanguageServerTaskFn) -> bool {
+    self.task_tx.send(task_fn).is_ok()
+  }
+
+  /// Panics if called more than once.
+  fn start(&mut self, ls: LanguageServer) {
+    let mut task_rx = self.task_rx.take().unwrap();
+    spawn(async move {
+      while let Some(task_fn) = task_rx.recv().await {
+        task_fn(ls.clone());
+      }
+    });
+  }
+}
+
 #[derive(Debug)]
 pub struct Inner {
   /// Cached versions of "fixed" assets that can either be inlined in Rust or
@@ -196,6 +238,7 @@ pub struct Inner {
   /// on disk or "open" within the client.
   pub documents: Documents,
   http_client: Arc<HttpClient>,
+  task_queue: LanguageServerTaskQueue,
   /// Handles module registries, which allow discovery of modules
   module_registries: ModuleRegistry,
   /// The path to the module registries cache
@@ -500,6 +543,7 @@ impl Inner {
       maybe_import_map_uri: None,
       maybe_package_json: None,
       fmt_options: Default::default(),
+      task_queue: Default::default(),
       lint_options: Default::default(),
       maybe_testing_server: None,
       module_registries,
@@ -1023,6 +1067,41 @@ impl Inner {
       self.lint_options = lint_options;
       self.fmt_options = fmt_options;
       self.recreate_http_client_and_dependents().await?;
+      if let Some(config_file) = self.config.maybe_config_file() {
+        if let Ok((compiler_options, _)) = config_file.to_compiler_options() {
+          if let Some(compiler_options_obj) = compiler_options.as_object() {
+            if let Some(jsx_import_source) =
+              compiler_options_obj.get("jsxImportSource")
+            {
+              if let Some(jsx_import_source) = jsx_import_source.as_str() {
+                let cache_params = lsp_custom::CacheParams {
+                  referrer: TextDocumentIdentifier {
+                    uri: config_file.specifier.clone(),
+                  },
+                  uris: vec![TextDocumentIdentifier {
+                    uri: Url::parse(&format!(
+                      "data:application/typescript;base64,{}",
+                      base64::engine::general_purpose::STANDARD.encode(
+                        format!("import '{jsx_import_source}/jsx-runtime';")
+                      )
+                    ))
+                    .unwrap(),
+                  }],
+                };
+                self.task_queue.queue_task(Box::new(|ls: LanguageServer| {
+                  spawn(async move {
+                    if let Err(err) =
+                      ls.cache_request(Some(json!(cache_params))).await
+                    {
+                      lsp_warn!("{}", err);
+                    }
+                  });
+                }));
+              }
+            }
+          }
+        }
+      }
     }
 
     Ok(())
@@ -3257,9 +3336,8 @@ impl tower_lsp::LanguageServer for LanguageServer {
       ls.refresh_documents_config().await;
       ls.diagnostics_server.invalidate_all();
       ls.send_diagnostics_update();
-    }
-
-    lsp_log!("Server ready.");
+      ls.task_queue.start(self.clone());
+    };
 
     if upgrade_check_enabled() {
       // spawn to avoid lsp send/sync requirement, but also just
@@ -3282,6 +3360,8 @@ impl tower_lsp::LanguageServer for LanguageServer {
         }
       });
     }
+
+    lsp_log!("Server ready.");
   }
 
   async fn shutdown(&self) -> LspResult<()> {
@@ -3596,10 +3676,6 @@ impl Inner {
     let referrer = self
       .url_map
       .normalize_url(&params.referrer.uri, LspUrlKind::File);
-    if !self.is_diagnosable(&referrer) {
-      return Ok(None);
-    }
-
     let mark = self.performance.mark_with_args("lsp.cache", &params);
     let roots = if !params.uris.is_empty() {
       params
diff --git a/cli/tests/integration/lsp_tests.rs b/cli/tests/integration/lsp_tests.rs
index 78aff93ab..e9a4db535 100644
--- a/cli/tests/integration/lsp_tests.rs
+++ b/cli/tests/integration/lsp_tests.rs
@@ -9469,6 +9469,61 @@ export function B() {
   client.shutdown();
 }
 
+#[test]
+fn lsp_jsx_import_source_config_file_automatic_cache() {
+  let context = TestContextBuilder::new()
+    .use_http_server()
+    .use_temp_cwd()
+    .build();
+  let temp_dir = context.temp_dir();
+  temp_dir.write(
+    "deno.json",
+    json!({
+      "compilerOptions": {
+        "jsx": "react-jsx",
+        "jsxImportSource": "http://localhost:4545/jsx",
+      },
+    })
+    .to_string(),
+  );
+  let mut client = context.new_lsp_command().build();
+  client.initialize_default();
+  let mut diagnostics = client.did_open(json!({
+    "textDocument": {
+      "uri": temp_dir.uri().join("file.tsx").unwrap(),
+      "languageId": "typescriptreact",
+      "version": 1,
+      "text": "
+        export function Foo() {
+          return <div></div>;
+        }
+      ",
+    },
+  }));
+  // The caching is done on an asynchronous task spawned after init, so there's
+  // a chance it wasn't done in time and we need to wait for another batch of
+  // diagnostics.
+  while !diagnostics.all().is_empty() {
+    std::thread::sleep(std::time::Duration::from_millis(50));
+    // The post-cache diagnostics update triggers inconsistently on CI for some
+    // reason. Force it with this notification.
+    diagnostics = client.did_open(json!({
+      "textDocument": {
+        "uri": temp_dir.uri().join("file.tsx").unwrap(),
+        "languageId": "typescriptreact",
+        "version": 1,
+        "text": "
+          export function Foo() {
+            return <div></div>;
+          }
+        ",
+      },
+    }));
+  }
+  assert_eq!(diagnostics.all(), vec![]);
+  client.shutdown();
+}
+
 #[derive(Debug, Clone, Deserialize, PartialEq)]
 #[serde(rename_all = "camelCase")]
 struct TestData {