From 6b7107c81e150ccddbb6d1b23f05ea20fc8ca20f Mon Sep 17 00:00:00 2001
From: Yang Yu <xuyang715812@icloud.com>
Date: Thu, 16 Jul 2026 03:17:24 +0800
Subject: [PATCH] Fix first auto-import completion for unloaded modules

Module completion only searched modules already loaded into the transaction, so discoverable stdlib and site-package modules were absent until an explicit import loaded them. Merge import-prefix candidates from the active file's configuration into the fuzzy module search and deduplicate them. Add a regression test for completing json before it has been imported.
---
 pyrefly/lib/lsp/wasm/completion.rs |  2 +-
 pyrefly/lib/state/lsp.rs           | 10 ++++++---
 pyrefly/lib/test/lsp/completion.rs | 35 ++++++++++++++++++++++++++++++
 3 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/pyrefly/lib/lsp/wasm/completion.rs b/pyrefly/lib/lsp/wasm/completion.rs
index 8646055f06..8e4f6a7f2c 100644
--- a/pyrefly/lib/lsp/wasm/completion.rs
+++ b/pyrefly/lib/lsp/wasm/completion.rs
@@ -707,7 +707,7 @@ impl Transaction<'_> {
                 });
             }
 
-            for module_name in self.search_modules_fuzzy(identifier_text) {
+            for module_name in self.search_modules_fuzzy(handle, identifier_text) {
                 if module_name == handle.module() {
                     continue;
                 }
diff --git a/pyrefly/lib/state/lsp.rs b/pyrefly/lib/state/lsp.rs
index 7628b76fda..9dd210337c 100644
--- a/pyrefly/lib/state/lsp.rs
+++ b/pyrefly/lib/state/lsp.rs
@@ -2917,11 +2917,15 @@ impl<'a> Transaction<'a> {
         .map(|item| TextRangeWithModule::new(item.module, item.definition_range))
     }
 
-    pub(crate) fn search_modules_fuzzy(&self, pattern: &str) -> Vec<ModuleName> {
+    pub(crate) fn search_modules_fuzzy(&self, handle: &Handle, pattern: &str) -> Vec<ModuleName> {
         let matcher = SkimMatcherV2::default().smart_case();
         let mut results = Vec::new();
+        // `self.modules()` only contains modules that have already been loaded. Include modules
+        // discoverable from the active file's import paths so auto-import works on the first try.
+        let mut module_names: HashSet<_> = self.modules().into_iter().collect();
+        module_names.extend(self.import_prefixes(handle, ModuleName::from_str(pattern)));
 
-        for module_name in self.modules() {
+        for module_name in module_names {
             let module_name_str = module_name.as_str();
 
             // Skip builtins module
@@ -3010,7 +3014,7 @@ impl<'a> Transaction<'a> {
                         &mut import_actions,
                         unknown_name,
                     );
-                    for module_name in self.search_modules_fuzzy(unknown_name) {
+                    for module_name in self.search_modules_fuzzy(handle, unknown_name) {
                         if module_name == handle.module() {
                             continue;
                         }
diff --git a/pyrefly/lib/test/lsp/completion.rs b/pyrefly/lib/test/lsp/completion.rs
index 5aa6eeda86..206d322e6e 100644
--- a/pyrefly/lib/test/lsp/completion.rs
+++ b/pyrefly/lib/test/lsp/completion.rs
@@ -2398,6 +2398,41 @@ jso
     );
 }
 
+#[test]
+fn autoimport_suggests_unloaded_stdlib_module() {
+    let code = r#"
+jso
+#  ^
+"#;
+    let mut env = TestEnv::new().with_default_require_level(Require::Exports);
+    env.add("main", code);
+    let (state, handle_for) = env.to_state();
+    let handle = handle_for("main");
+    let position = extract_cursors_for_test(code)[0];
+
+    let completions =
+        state
+            .transaction()
+            .completion(&handle, position, ImportFormat::Absolute, true, None);
+    let json = completions
+        .iter()
+        .find(|item| item.label == "json")
+        .expect("expected an auto-import completion for the unloaded json module");
+    assert!(
+        json.detail
+            .as_ref()
+            .is_some_and(|detail| detail.contains("import json")),
+        "expected json completion to add an import, got {:?}",
+        json.detail
+    );
+    assert!(
+        json.additional_text_edits
+            .as_ref()
+            .is_some_and(|edits| !edits.is_empty()),
+        "expected json completion to include an import edit"
+    );
+}
+
 #[test]
 fn autoimport_does_not_duplicate_existing_module_import_after_another_import() {
     let code = r#"