fix(lsp): rewrite import for 'infer return type' action (#24685)

3 files changed, 22 insertions, 6 deletions

diff --git a/cli/lsp/analysis.rs b/cli/lsp/analysis.rs
index 8480f6e1d..a6a69cbf0 100644
--- a/cli/lsp/analysis.rs
+++ b/cli/lsp/analysis.rs
@@ -73,8 +73,9 @@ static PREFERRED_FIXES: Lazy<HashMap<&'static str, (u32, bool)>> =
     .collect()
   });
 
-static IMPORT_SPECIFIER_RE: Lazy<Regex> =
-  lazy_regex::lazy_regex!(r#"\sfrom\s+["']([^"']*)["']"#);
+static IMPORT_SPECIFIER_RE: Lazy<Regex> = lazy_regex::lazy_regex!(
+  r#"\sfrom\s+["']([^"']*)["']|import\s*\(\s*["']([^"']*)["']\s*\)"#
+);
 
 const SUPPORTED_EXTENSIONS: &[&str] = &[
   ".ts", ".tsx", ".js", ".jsx", ".mjs", ".mts", ".cjs", ".cts", ".d.ts",
@@ -528,7 +529,8 @@ pub fn fix_ts_import_changes(
         .map(|line| {
           // This assumes that there's only one import per line.
           if let Some(captures) = IMPORT_SPECIFIER_RE.captures(line) {
-            let specifier = captures.get(1).unwrap().as_str();
+            let specifier =
+              captures.iter().skip(1).find_map(|s| s).unwrap().as_str();
             if let Some(new_specifier) =
               import_mapper.check_unresolved_specifier(specifier, referrer)
             {
diff --git a/cli/lsp/language_server.rs b/cli/lsp/language_server.rs
index aa1b4c282..9680c63f9 100644
--- a/cli/lsp/language_server.rs
+++ b/cli/lsp/language_server.rs
@@ -1810,7 +1810,10 @@ impl Inner {
         LspError::internal_error()
       })?;
       code_action
-    } else if kind.as_str().starts_with(CodeActionKind::REFACTOR.as_str()) {
+    } else if let Some(kind_suffix) = kind
+      .as_str()
+      .strip_prefix(CodeActionKind::REFACTOR.as_str())
+    {
       let mut code_action = params;
       let action_data: refactor::RefactorCodeActionData = from_value(data)
         .map_err(|err| {
@@ -1819,7 +1822,7 @@ impl Inner {
         })?;
       let asset_or_doc = self.get_asset_or_document(&action_data.specifier)?;
       let line_index = asset_or_doc.line_index();
-      let refactor_edit_info = self
+      let mut refactor_edit_info = self
         .ts_server
         .get_edits_for_refactor(
           self.snapshot(),
@@ -1841,6 +1844,17 @@ impl Inner {
           asset_or_doc.scope().cloned(),
         )
         .await?;
+      if kind_suffix == ".rewrite.function.returnType" {
+        refactor_edit_info.edits = fix_ts_import_changes(
+          &action_data.specifier,
+          &refactor_edit_info.edits,
+          &self.get_ts_response_import_mapper(&action_data.specifier),
+        )
+        .map_err(|err| {
+          error!("Unable to remap changes: {:#}", err);
+          LspError::internal_error()
+        })?
+      }
       code_action.edit = refactor_edit_info.to_workspace_edit(self)?;
       code_action
     } else {
diff --git a/cli/lsp/tsc.rs b/cli/lsp/tsc.rs
index 6c808a5e4..4592edad4 100644
--- a/cli/lsp/tsc.rs
+++ b/cli/lsp/tsc.rs
@@ -2944,7 +2944,7 @@ pub fn file_text_changes_to_workspace_edit(
 #[derive(Debug, Deserialize, Serialize)]
 #[serde(rename_all = "camelCase")]
 pub struct RefactorEditInfo {
-  edits: Vec<FileTextChanges>,
+  pub edits: Vec<FileTextChanges>,
   #[serde(skip_serializing_if = "Option::is_none")]
   pub rename_location: Option<u32>,
 }