[ty] Add "import ..." code-action for unresolved references (#21629)

## Summary

Originally I planned to feed this in as a `fix` but I realized that we
probably don't want to be trying to resolve import suggestions while
we're doing type inference. Thus I implemented this as a fallback when
there's no fixes on a diagnostic, which can use the full lsp machinery.

Fixes https://github.com/astral-sh/ty/issues/1552

## Test Plan

Works in the IDE, added some e2e tests.

crates/ty_ide/src/all_symbols.rs

@@ -143,7 +143,7 @@ ABCDEFGHIJKLMNOP = 'https://api.example.com'
 
     impl CursorTest {
         fn all_symbols(&self, query: &str) -> String {
-            let symbols = all_symbols(&self.db, &QueryPattern::new(query));
+            let symbols = all_symbols(&self.db, &QueryPattern::fuzzy(query));
 
             if symbols.is_empty() {
                 return "No symbols found".to_string();

crates/ty_ide/src/code_action.rs

@@ -0,0 +1,43 @@
+use crate::{completion, find_node::covering_node};
+use ruff_db::{files::File, parsed::parsed_module};
+use ruff_diagnostics::Edit;
+use ruff_text_size::TextRange;
+use ty_project::Db;
+use ty_python_semantic::types::UNRESOLVED_REFERENCE;
+
+/// A `QuickFix` Code Action
+#[derive(Debug, Clone)]
+pub struct QuickFix {
+    pub title: String,
+    pub edits: Vec<Edit>,
+    pub preferred: bool,
+}
+
+pub fn code_actions(
+    db: &dyn Db,
+    file: File,
+    diagnostic_range: TextRange,
+    diagnostic_id: &str,
+) -> Option<Vec<QuickFix>> {
+    let registry = db.lint_registry();
+    let Ok(lint_id) = registry.get(diagnostic_id) else {
+        return None;
+    };
+    if lint_id.name() == UNRESOLVED_REFERENCE.name() {
+        let parsed = parsed_module(db, file).load(db);
+        let node = covering_node(parsed.syntax().into(), diagnostic_range).node();
+        let symbol = &node.expr_name()?.id;
+
+        let fixes = completion::missing_imports(db, file, &parsed, symbol, node)
+            .into_iter()
+            .map(|import| QuickFix {
+                title: import.label,
+                edits: vec![import.edit],
+                preferred: true,
+            })
+            .collect();
+        Some(fixes)
+    } else {
+        None
+    }
+}

crates/ty_ide/src/completion.rs

@@ -4,8 +4,8 @@ use ruff_db::files::File;
 use ruff_db::parsed::{ParsedModuleRef, parsed_module};
 use ruff_db::source::source_text;
 use ruff_diagnostics::Edit;
-use ruff_python_ast as ast;
 use ruff_python_ast::name::Name;
+use ruff_python_ast::{self as ast, AnyNodeRef};
 use ruff_python_codegen::Stylist;
 use ruff_python_parser::{Token, TokenAt, TokenKind, Tokens};
 use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
@@ -37,9 +37,9 @@ impl<'db> Completions<'db> {
     /// the user has typed as part of the next symbol they are writing.
     /// This collection will treat it as a query when present, and only
     /// add completions that match it.
-    fn new(db: &'db dyn Db, typed: Option<&str>) -> Completions<'db> {
+    fn fuzzy(db: &'db dyn Db, typed: Option<&str>) -> Completions<'db> {
         let query = typed
-            .map(QueryPattern::new)
+            .map(QueryPattern::fuzzy)
             .unwrap_or_else(QueryPattern::matches_all_symbols);
         Completions {
             db,
@@ -48,6 +48,15 @@ impl<'db> Completions<'db> {
         }
     }
 
+    fn exactly(db: &'db dyn Db, symbol: &str) -> Completions<'db> {
+        let query = QueryPattern::exactly(symbol);
+        Completions {
+            db,
+            items: vec![],
+            query,
+        }
+    }
+
     /// Convert this collection into a simple
     /// sequence of completions.
     fn into_completions(mut self) -> Vec<Completion<'db>> {
@@ -57,6 +66,21 @@ impl<'db> Completions<'db> {
         self.items
     }
 
+    fn into_imports(mut self) -> Vec<ImportEdit> {
+        self.items.sort_by(compare_suggestions);
+        self.items
+            .dedup_by(|c1, c2| (&c1.name, c1.module_name) == (&c2.name, c2.module_name));
+        self.items
+            .into_iter()
+            .filter_map(|item| {
+                Some(ImportEdit {
+                    label: format!("import {}.{}", item.module_name?, item.name),
+                    edit: item.import?,
+                })
+            })
+            .collect()
+    }
+
     /// Attempts to adds the given completion to this collection.
     ///
     /// When added, `true` is returned.
@@ -369,7 +393,7 @@ pub fn completion<'db>(
         return vec![];
     }
 
-    let mut completions = Completions::new(db, typed.as_deref());
+    let mut completions = Completions::fuzzy(db, typed.as_deref());
 
     if let Some(import) = ImportStatement::detect(db, file, &parsed, tokens, typed.as_deref()) {
         import.add_completions(db, file, &mut completions);
@@ -417,6 +441,25 @@ pub fn completion<'db>(
     completions.into_completions()
 }
 
+pub(crate) struct ImportEdit {
+    pub label: String,
+    pub edit: Edit,
+}
+
+pub(crate) fn missing_imports(
+    db: &dyn Db,
+    file: File,
+    parsed: &ParsedModuleRef,
+    symbol: &str,
+    node: AnyNodeRef,
+) -> Vec<ImportEdit> {
+    let mut completions = Completions::exactly(db, symbol);
+    let scoped = ScopedTarget { node };
+    add_unimported_completions(db, file, parsed, scoped, &mut completions);
+
+    completions.into_imports()
+}
+
 /// Adds completions derived from keywords.
 ///
 /// This should generally only be used when offering "scoped" completions.

crates/ty_ide/src/lib.rs

@@ -3,6 +3,7 @@
     reason = "Prefer System trait methods over std methods in ty crates"
 )]
 mod all_symbols;
+mod code_action;
 mod completion;
 mod doc_highlights;
 mod docstring;
@@ -27,6 +28,7 @@ mod symbols;
 mod workspace_symbols;
 
 pub use all_symbols::{AllSymbolInfo, all_symbols};
+pub use code_action::{QuickFix, code_actions};
 pub use completion::{Completion, CompletionKind, CompletionSettings, completion};
 pub use doc_highlights::document_highlights;
 pub use document_symbols::document_symbols;

crates/ty_ide/src/symbols.rs

@@ -23,11 +23,12 @@ use crate::completion::CompletionKind;
 pub struct QueryPattern {
     re: Option<Regex>,
     original: String,
+    original_is_exact: bool,
 }
 
 impl QueryPattern {
     /// Create a new query pattern from a literal search string given.
-    pub fn new(literal_query_string: &str) -> QueryPattern {
+    pub fn fuzzy(literal_query_string: &str) -> QueryPattern {
         let mut pattern = "(?i)".to_string();
         for ch in literal_query_string.chars() {
             pattern.push_str(&regex::escape(ch.encode_utf8(&mut [0; 4])));
@@ -41,6 +42,16 @@ impl QueryPattern {
         QueryPattern {
             re: Regex::new(&pattern).ok(),
             original: literal_query_string.to_string(),
+            original_is_exact: false,
+        }
+    }
+
+    /// Create a new query
+    pub fn exactly(symbol: &str) -> QueryPattern {
+        QueryPattern {
+            re: None,
+            original: symbol.to_string(),
+            original_is_exact: true,
         }
     }
 
@@ -49,6 +60,7 @@ impl QueryPattern {
         QueryPattern {
             re: None,
             original: String::new(),
+            original_is_exact: false,
         }
     }
 
@@ -59,6 +71,8 @@ impl QueryPattern {
     pub fn is_match_symbol_name(&self, symbol_name: &str) -> bool {
         if let Some(ref re) = self.re {
             re.is_match(symbol_name)
+        } else if self.original_is_exact {
+            symbol_name == self.original
         } else {
             // This is a degenerate case. The only way
             // we should get here is if the query string
@@ -75,13 +89,13 @@ impl QueryPattern {
     /// incorrectly. That is, it's possible that this query will match all
     /// inputs but this still returns `false`.
     pub fn will_match_everything(&self) -> bool {
-        self.re.is_none()
+        self.re.is_none() && self.original.is_empty()
     }
 }
 
 impl From<&str> for QueryPattern {
     fn from(literal_query_string: &str) -> QueryPattern {
-        QueryPattern::new(literal_query_string)
+        QueryPattern::fuzzy(literal_query_string)
     }
 }
 
@@ -565,7 +579,7 @@ impl SourceOrderVisitor<'_> for SymbolVisitor {
 #[cfg(test)]
 mod tests {
     fn matches(query: &str, symbol: &str) -> bool {
-        super::QueryPattern::new(query).is_match_symbol_name(symbol)
+        super::QueryPattern::fuzzy(query).is_match_symbol_name(symbol)
     }
 
     #[test]

crates/ty_ide/src/workspace_symbols.rs

@@ -12,7 +12,7 @@ pub fn workspace_symbols(db: &dyn Db, query: &str) -> Vec<WorkspaceSymbolInfo> {
 
     let project = db.project();
 
-    let query = QueryPattern::new(query);
+    let query = QueryPattern::fuzzy(query);
     let files = project.files(db);
     let results = std::sync::Mutex::new(Vec::new());
     {

crates/ty_python_semantic/src/types.rs

@@ -24,7 +24,7 @@ pub(crate) use self::builder::{IntersectionBuilder, UnionBuilder};
 pub use self::cyclic::CycleDetector;
 pub(crate) use self::cyclic::{PairVisitor, TypeTransformer};
 pub(crate) use self::diagnostic::register_lints;
-pub use self::diagnostic::{TypeCheckDiagnostics, UNDEFINED_REVEAL};
+pub use self::diagnostic::{TypeCheckDiagnostics, UNDEFINED_REVEAL, UNRESOLVED_REFERENCE};
 pub(crate) use self::infer::{
     TypeContext, infer_deferred_types, infer_definition_types, infer_expression_type,
     infer_expression_types, infer_scope_types, static_expression_truthiness,

crates/ty_python_semantic/src/types/diagnostic.rs

@@ -1844,7 +1844,7 @@ declare_lint! {
     /// ```python
     /// print(x)  # NameError: name 'x' is not defined
     /// ```
-    pub(crate) static UNRESOLVED_REFERENCE = {
+    pub static UNRESOLVED_REFERENCE = {
         summary: "detects references to names that are not defined",
         status: LintStatus::stable("0.0.1-alpha.1"),
         default_level: Level::Error,

crates/ty_server/src/server/api/requests/code_action.rs

@@ -1,16 +1,23 @@
 use std::borrow::Cow;
+use std::collections::HashMap;
 
-use lsp_types::{self as types, Url, request as req};
+use lsp_types::{self as types, NumberOrString, TextEdit, Url, request as req};
+use ruff_db::files::File;
+use ruff_diagnostics::Edit;
+use ruff_text_size::Ranged;
+use ty_ide::code_actions;
 use ty_project::ProjectDatabase;
 use types::{CodeActionKind, CodeActionOrCommand};
 
-use crate::DIAGNOSTIC_NAME;
+use crate::db::Db;
+use crate::document::{RangeExt, ToRangeExt};
 use crate::server::Result;
 use crate::server::api::RequestHandler;
 use crate::server::api::diagnostics::DiagnosticData;
 use crate::server::api::traits::{BackgroundDocumentRequestHandler, RetriableRequestHandler};
 use crate::session::DocumentSnapshot;
 use crate::session::client::Client;
+use crate::{DIAGNOSTIC_NAME, PositionEncoding};
 
 pub(crate) struct CodeActionRequestHandler;
 
@@ -24,23 +31,24 @@ impl BackgroundDocumentRequestHandler for CodeActionRequestHandler {
     }
 
     fn run_with_snapshot(
-        _db: &ProjectDatabase,
+        db: &ProjectDatabase,
-        _snapshot: &DocumentSnapshot,
+        snapshot: &DocumentSnapshot,
         _client: &Client,
         params: types::CodeActionParams,
     ) -> Result<Option<types::CodeActionResponse>> {
         let diagnostics = params.context.diagnostics;
 
+        let Some(file) = snapshot.to_notebook_or_file(db) else {
+            return Ok(None);
+        };
         let mut actions = Vec::new();
 
         for mut diagnostic in diagnostics.into_iter().filter(|diagnostic| {
             diagnostic.source.as_deref() == Some(DIAGNOSTIC_NAME)
                 && range_intersect(&diagnostic.range, &params.range)
         }) {
-            let Some(data) = diagnostic.data.take() else {
+            // If the diagnostic includes fixes, offer those up as options.
-                continue;
+            if let Some(data) = diagnostic.data.take() {
-            };
-
                 let data: DiagnosticData = match serde_json::from_value(data) {
                     Ok(data) => data,
                     Err(err) => {
@@ -52,7 +60,7 @@ impl BackgroundDocumentRequestHandler for CodeActionRequestHandler {
                 actions.push(CodeActionOrCommand::CodeAction(lsp_types::CodeAction {
                     title: data.fix_title,
                     kind: Some(CodeActionKind::QUICKFIX),
-                diagnostics: Some(vec![diagnostic]),
+                    diagnostics: Some(vec![diagnostic.clone()]),
                     edit: Some(lsp_types::WorkspaceEdit {
                         changes: Some(data.edits),
                         document_changes: None,
@@ -65,13 +73,69 @@ impl BackgroundDocumentRequestHandler for CodeActionRequestHandler {
                 }));
             }
 
-        if actions.is_empty() {
+            // Try to find other applicable actions.
-            return Ok(None);
+            //
+            // This is only for actions that are messy to compute at the time of the diagnostic.
+            // For instance, suggesting imports requires finding symbols for the entire project,
+            // which is dubious when you're in the middle of resolving symbols.
+            let url = snapshot.url();
+            let encoding = snapshot.encoding();
+            if let Some(NumberOrString::String(diagnostic_id)) = &diagnostic.code
+                && let Some(range) = diagnostic.range.to_text_range(db, file, url, encoding)
+                && let Some(fixes) = code_actions(db, file, range, diagnostic_id)
+            {
+                for action in fixes {
+                    actions.push(CodeActionOrCommand::CodeAction(lsp_types::CodeAction {
+                        title: action.title,
+                        kind: Some(CodeActionKind::QUICKFIX),
+                        diagnostics: Some(vec![diagnostic.clone()]),
+                        edit: Some(lsp_types::WorkspaceEdit {
+                            changes: to_lsp_edits(db, file, encoding, action.edits),
+                            document_changes: None,
+                            change_annotations: None,
+                        }),
+                        is_preferred: Some(action.preferred),
+                        command: None,
+                        disabled: None,
+                        data: None,
+                    }));
+                }
+            }
         }
 
+        if actions.is_empty() {
+            Ok(None)
+        } else {
             Ok(Some(actions))
         }
     }
+}
+
+fn to_lsp_edits(
+    db: &dyn Db,
+    file: File,
+    encoding: PositionEncoding,
+    edits: Vec<Edit>,
+) -> Option<HashMap<Url, Vec<TextEdit>>> {
+    let mut lsp_edits: HashMap<Url, Vec<lsp_types::TextEdit>> = HashMap::new();
+
+    for edit in edits {
+        let location = edit
+            .range()
+            .to_lsp_range(db, file, encoding)?
+            .to_location()?;
+
+        lsp_edits
+            .entry(location.uri)
+            .or_default()
+            .push(lsp_types::TextEdit {
+                range: location.range,
+                new_text: edit.content().unwrap_or_default().to_string(),
+            });
+    }
+
+    Some(lsp_edits)
+}
 
 fn range_intersect(range: &lsp_types::Range, other: &lsp_types::Range) -> bool {
     let start = range.start.max(other.start);

crates/ty_server/tests/e2e/code_actions.rs

@@ -1,8 +1,47 @@
+use crate::{TestServer, TestServerBuilder};
 use anyhow::Result;
-use lsp_types::{Position, Range, request::CodeActionRequest};
+use lsp_types::{DocumentDiagnosticReportResult, Position, Range, request::CodeActionRequest};
 use ruff_db::system::SystemPath;
 
-use crate::TestServerBuilder;
+fn code_actions_at(
+    server: &TestServer,
+    diagnostics: DocumentDiagnosticReportResult,
+    file: &SystemPath,
+    range: Range,
+) -> lsp_types::CodeActionParams {
+    lsp_types::CodeActionParams {
+        text_document: lsp_types::TextDocumentIdentifier {
+            uri: server.file_uri(file),
+        },
+        range,
+        context: lsp_types::CodeActionContext {
+            diagnostics: match diagnostics {
+                lsp_types::DocumentDiagnosticReportResult::Report(
+                    lsp_types::DocumentDiagnosticReport::Full(report),
+                ) => report.full_document_diagnostic_report.items,
+                _ => panic!("Expected full diagnostic report"),
+            },
+            only: None,
+            trigger_kind: None,
+        },
+        work_done_progress_params: lsp_types::WorkDoneProgressParams::default(),
+        partial_result_params: lsp_types::PartialResultParams::default(),
+    }
+}
+
+#[allow(clippy::cast_possible_truncation)]
+fn full_range(input: &str) -> Range {
+    let (num_lines, last_line) = input
+        .lines()
+        .enumerate()
+        .last()
+        .expect("non-empty document");
+    let last_char = last_line.len() as u32;
+    Range::new(
+        Position::new(0, 0),
+        Position::new(num_lines as u32, last_char),
+    )
+}
 
 #[test]
 fn code_action() -> Result<()> {
@@ -30,27 +69,10 @@ unused-ignore-comment = \"warn\"
 
     // Wait for diagnostics to be computed.
     let diagnostics = server.document_diagnostic_request(foo, None);
+    let range = full_range(foo_content);
+    let code_action_params = code_actions_at(&server, diagnostics, foo, range);
 
     // Get code actions for the line with the unused ignore comment.
-    let code_action_params = lsp_types::CodeActionParams {
-        text_document: lsp_types::TextDocumentIdentifier {
-            uri: server.file_uri(foo),
-        },
-        range: Range::new(Position::new(0, 0), Position::new(0, 43)),
-        context: lsp_types::CodeActionContext {
-            diagnostics: match diagnostics {
-                lsp_types::DocumentDiagnosticReportResult::Report(
-                    lsp_types::DocumentDiagnosticReport::Full(report),
-                ) => report.full_document_diagnostic_report.items,
-                _ => panic!("Expected full diagnostic report"),
-            },
-            only: None,
-            trigger_kind: None,
-        },
-        work_done_progress_params: lsp_types::WorkDoneProgressParams::default(),
-        partial_result_params: lsp_types::PartialResultParams::default(),
-    };
-
     let code_action_id = server.send_request::<CodeActionRequest>(code_action_params);
     let code_actions = server.await_response::<CodeActionRequest>(&code_action_id);
 
@@ -88,24 +110,8 @@ unused-ignore-comment = \"warn\"
 
     // Get code actions for a range that doesn't overlap with the diagnostic.
     // The diagnostic is at characters 12-42, so we request actions for characters 0-10.
-    let code_action_params = lsp_types::CodeActionParams {
+    let range = Range::new(Position::new(0, 0), Position::new(0, 10));
-        text_document: lsp_types::TextDocumentIdentifier {
+    let code_action_params = code_actions_at(&server, diagnostics, foo, range);
-            uri: server.file_uri(foo),
-        },
-        range: Range::new(Position::new(0, 0), Position::new(0, 10)),
-        context: lsp_types::CodeActionContext {
-            diagnostics: match diagnostics {
-                lsp_types::DocumentDiagnosticReportResult::Report(
-                    lsp_types::DocumentDiagnosticReport::Full(report),
-                ) => report.full_document_diagnostic_report.items,
-                _ => panic!("Expected full diagnostic report"),
-            },
-            only: None,
-            trigger_kind: None,
-        },
-        work_done_progress_params: lsp_types::WorkDoneProgressParams::default(),
-        partial_result_params: lsp_types::PartialResultParams::default(),
-    };
 
     let code_action_id = server.send_request::<CodeActionRequest>(code_action_params);
     let code_actions = server.await_response::<CodeActionRequest>(&code_action_id);
@@ -115,3 +121,152 @@ unused-ignore-comment = \"warn\"
 
     Ok(())
 }
+
+// `Literal` is available from two places so we should suggest two possible imports
+#[test]
+fn code_action_undefined_reference_multi() -> Result<()> {
+    let workspace_root = SystemPath::new("src");
+    let foo = SystemPath::new("src/foo.py");
+    let foo_content = "\
+x: Literal[1] = 1
+";
+
+    let ty_toml = SystemPath::new("ty.toml");
+    let ty_toml_content = "\
+[rules]
+unused-ignore-comment = \"warn\"
+";
+
+    let mut server = TestServerBuilder::new()?
+        .with_workspace(workspace_root, None)?
+        .with_file(ty_toml, ty_toml_content)?
+        .with_file(foo, foo_content)?
+        .enable_pull_diagnostics(true)
+        .build()
+        .wait_until_workspaces_are_initialized();
+
+    server.open_text_document(foo, &foo_content, 1);
+
+    // Wait for diagnostics to be computed.
+    let diagnostics = server.document_diagnostic_request(foo, None);
+    let range = full_range(foo_content);
+    let code_action_params = code_actions_at(&server, diagnostics, foo, range);
+
+    // Get code actions
+    let code_action_id = server.send_request::<CodeActionRequest>(code_action_params);
+    let code_actions = server.await_response::<CodeActionRequest>(&code_action_id);
+
+    insta::assert_json_snapshot!(code_actions);
+
+    Ok(())
+}
+
+// Using an unimported decorator `@deprecated`
+#[test]
+fn code_action_undefined_decorator() -> Result<()> {
+    let workspace_root = SystemPath::new("src");
+    let foo = SystemPath::new("src/foo.py");
+    let foo_content = r#"\
+@deprecated("do not use!!!")
+def my_func(): ...
+"#;
+
+    let ty_toml = SystemPath::new("ty.toml");
+    let ty_toml_content = "";
+
+    let mut server = TestServerBuilder::new()?
+        .with_workspace(workspace_root, None)?
+        .with_file(ty_toml, ty_toml_content)?
+        .with_file(foo, foo_content)?
+        .enable_pull_diagnostics(true)
+        .build()
+        .wait_until_workspaces_are_initialized();
+
+    server.open_text_document(foo, &foo_content, 1);
+
+    // Wait for diagnostics to be computed.
+    let diagnostics = server.document_diagnostic_request(foo, None);
+    let range = full_range(foo_content);
+    let code_action_params = code_actions_at(&server, diagnostics, foo, range);
+
+    // Get code actions
+    let code_action_id = server.send_request::<CodeActionRequest>(code_action_params);
+    let code_actions = server.await_response::<CodeActionRequest>(&code_action_id);
+
+    insta::assert_json_snapshot!(code_actions);
+
+    Ok(())
+}
+
+// Accessing `typing.Literal` without `typing` imported (ideally we suggest importing `typing`)
+#[test]
+fn code_action_attribute_access_on_unimported() -> Result<()> {
+    let workspace_root = SystemPath::new("src");
+    let foo = SystemPath::new("src/foo.py");
+    let foo_content = "\
+x: typing.Literal[1] = 1
+";
+
+    let ty_toml = SystemPath::new("ty.toml");
+    let ty_toml_content = "";
+
+    let mut server = TestServerBuilder::new()?
+        .with_workspace(workspace_root, None)?
+        .with_file(ty_toml, ty_toml_content)?
+        .with_file(foo, foo_content)?
+        .enable_pull_diagnostics(true)
+        .build()
+        .wait_until_workspaces_are_initialized();
+
+    server.open_text_document(foo, &foo_content, 1);
+
+    // Wait for diagnostics to be computed.
+    let diagnostics = server.document_diagnostic_request(foo, None);
+    let range = full_range(foo_content);
+    let code_action_params = code_actions_at(&server, diagnostics, foo, range);
+
+    // Get code actions
+    let code_action_id = server.send_request::<CodeActionRequest>(code_action_params);
+    let code_actions = server.await_response::<CodeActionRequest>(&code_action_id);
+
+    insta::assert_json_snapshot!(code_actions);
+
+    Ok(())
+}
+
+// Accessing `html.parser` when we've imported `html` but not `html.parser`
+#[test]
+fn code_action_possible_missing_submodule_attribute() -> Result<()> {
+    let workspace_root = SystemPath::new("src");
+    let foo = SystemPath::new("src/foo.py");
+    let foo_content = "\
+import html
+html.parser
+";
+
+    let ty_toml = SystemPath::new("ty.toml");
+    let ty_toml_content = "";
+
+    let mut server = TestServerBuilder::new()?
+        .with_workspace(workspace_root, None)?
+        .with_file(ty_toml, ty_toml_content)?
+        .with_file(foo, foo_content)?
+        .enable_pull_diagnostics(true)
+        .build()
+        .wait_until_workspaces_are_initialized();
+
+    server.open_text_document(foo, &foo_content, 1);
+
+    // Wait for diagnostics to be computed.
+    let diagnostics = server.document_diagnostic_request(foo, None);
+    let range = full_range(foo_content);
+    let code_action_params = code_actions_at(&server, diagnostics, foo, range);
+
+    // Get code actions
+    let code_action_id = server.send_request::<CodeActionRequest>(code_action_params);
+    let code_actions = server.await_response::<CodeActionRequest>(&code_action_id);
+
+    insta::assert_json_snapshot!(code_actions);
+
+    Ok(())
+}

crates/ty_server/tests/e2e/snapshots/e2e__code_actions__code_action_attribute_access_on_unimported.snap

@@ -0,0 +1,5 @@
+---
+source: crates/ty_server/tests/e2e/code_actions.rs
+expression: code_actions
+---
+null

crates/ty_server/tests/e2e/snapshots/e2e__code_actions__code_action_possible_missing_submodule_attribute.snap

@@ -0,0 +1,5 @@
+---
+source: crates/ty_server/tests/e2e/code_actions.rs
+expression: code_actions
+---
+null

crates/ty_server/tests/e2e/snapshots/e2e__code_actions__code_action_undefined_decorator.snap

@@ -0,0 +1,98 @@
+---
+source: crates/ty_server/tests/e2e/code_actions.rs
+expression: code_actions
+---
+[
+  {
+    "title": "import typing_extensions.deprecated",
+    "kind": "quickfix",
+    "diagnostics": [
+      {
+        "range": {
+          "start": {
+            "line": 1,
+            "character": 1
+          },
+          "end": {
+            "line": 1,
+            "character": 11
+          }
+        },
+        "severity": 1,
+        "code": "unresolved-reference",
+        "codeDescription": {
+          "href": "https://ty.dev/rules#unresolved-reference"
+        },
+        "source": "ty",
+        "message": "Name `deprecated` used when not defined",
+        "relatedInformation": []
+      }
+    ],
+    "edit": {
+      "changes": {
+        "file://<temp_dir>/src/foo.py": [
+          {
+            "range": {
+              "start": {
+                "line": 0,
+                "character": 0
+              },
+              "end": {
+                "line": 0,
+                "character": 0
+              }
+            },
+            "newText": "from typing_extensions import deprecated\n"
+          }
+        ]
+      }
+    },
+    "isPreferred": true
+  },
+  {
+    "title": "import warnings.deprecated",
+    "kind": "quickfix",
+    "diagnostics": [
+      {
+        "range": {
+          "start": {
+            "line": 1,
+            "character": 1
+          },
+          "end": {
+            "line": 1,
+            "character": 11
+          }
+        },
+        "severity": 1,
+        "code": "unresolved-reference",
+        "codeDescription": {
+          "href": "https://ty.dev/rules#unresolved-reference"
+        },
+        "source": "ty",
+        "message": "Name `deprecated` used when not defined",
+        "relatedInformation": []
+      }
+    ],
+    "edit": {
+      "changes": {
+        "file://<temp_dir>/src/foo.py": [
+          {
+            "range": {
+              "start": {
+                "line": 0,
+                "character": 0
+              },
+              "end": {
+                "line": 0,
+                "character": 0
+              }
+            },
+            "newText": "from warnings import deprecated\n"
+          }
+        ]
+      }
+    },
+    "isPreferred": true
+  }
+]

crates/ty_server/tests/e2e/snapshots/e2e__code_actions__code_action_undefined_reference_multi.snap

@@ -0,0 +1,98 @@
+---
+source: crates/ty_server/tests/e2e/code_actions.rs
+expression: code_actions
+---
+[
+  {
+    "title": "import typing.Literal",
+    "kind": "quickfix",
+    "diagnostics": [
+      {
+        "range": {
+          "start": {
+            "line": 0,
+            "character": 3
+          },
+          "end": {
+            "line": 0,
+            "character": 10
+          }
+        },
+        "severity": 1,
+        "code": "unresolved-reference",
+        "codeDescription": {
+          "href": "https://ty.dev/rules#unresolved-reference"
+        },
+        "source": "ty",
+        "message": "Name `Literal` used when not defined",
+        "relatedInformation": []
+      }
+    ],
+    "edit": {
+      "changes": {
+        "file://<temp_dir>/src/foo.py": [
+          {
+            "range": {
+              "start": {
+                "line": 0,
+                "character": 0
+              },
+              "end": {
+                "line": 0,
+                "character": 0
+              }
+            },
+            "newText": "from typing import Literal\n"
+          }
+        ]
+      }
+    },
+    "isPreferred": true
+  },
+  {
+    "title": "import typing_extensions.Literal",
+    "kind": "quickfix",
+    "diagnostics": [
+      {
+        "range": {
+          "start": {
+            "line": 0,
+            "character": 3
+          },
+          "end": {
+            "line": 0,
+            "character": 10
+          }
+        },
+        "severity": 1,
+        "code": "unresolved-reference",
+        "codeDescription": {
+          "href": "https://ty.dev/rules#unresolved-reference"
+        },
+        "source": "ty",
+        "message": "Name `Literal` used when not defined",
+        "relatedInformation": []
+      }
+    ],
+    "edit": {
+      "changes": {
+        "file://<temp_dir>/src/foo.py": [
+          {
+            "range": {
+              "start": {
+                "line": 0,
+                "character": 0
+              },
+              "end": {
+                "line": 0,
+                "character": 0
+              }
+            },
+            "newText": "from typing_extensions import Literal\n"
+          }
+        ]
+      }
+    },
+    "isPreferred": true
+  }
+]

crates/ty_wasm/src/lib.rs

@@ -11,7 +11,7 @@ use ruff_db::system::{
     SystemPath, SystemPathBuf, SystemVirtualPath, WritableSystem,
 };
 use ruff_db::vendored::VendoredPath;
-use ruff_diagnostics::Applicability;
+use ruff_diagnostics::{Applicability, Edit};
 use ruff_notebook::Notebook;
 use ruff_python_formatter::formatted_file;
 use ruff_source_file::{LineIndex, OneIndexed, SourceLocation};
@@ -556,6 +556,51 @@ impl Workspace {
         Ok(result)
     }
 
+    #[wasm_bindgen(js_name = "codeActions")]
+    pub fn code_actions(
+        &self,
+        file_id: &FileHandle,
+        diagnostic: &Diagnostic,
+    ) -> Option<Vec<CodeAction>> {
+        // If the diagnostic includes fixes, offer those up as options.
+        let mut actions = Vec::new();
+        if let Some(action) = diagnostic.code_action(self) {
+            actions.push(action);
+        }
+
+        // Try to find other applicable actions.
+        //
+        // This is only for actions that are messy to compute at the time of the diagnostic.
+        // For instance, suggesting imports requires finding symbols for the entire project,
+        // which is dubious when you're in the middle of resolving symbols.
+        if let Some(range) = diagnostic.inner.range()
+            && let Some(fixes) = ty_ide::code_actions(
+                &self.db,
+                file_id.file,
+                range,
+                diagnostic.inner.id().as_str(),
+            )
+        {
+            for action in fixes {
+                actions.push(CodeAction {
+                    title: action.title,
+                    preferred: action.preferred,
+                    edits: action
+                        .edits
+                        .into_iter()
+                        .map(|edit| edit_to_text_edit(self, file_id.file, &edit))
+                        .collect(),
+                });
+            }
+        }
+
+        if actions.is_empty() {
+            None
+        } else {
+            Some(actions)
+        }
+    }
+
     #[wasm_bindgen(js_name = "signatureHelp")]
     pub fn signature_help(
         &self,
@@ -758,21 +803,10 @@ impl Diagnostic {
         let primary_span = self.inner.primary_span()?;
         let file = primary_span.expect_ty_file();
 
-        let source = source_text(&workspace.db, file);
-        let index = line_index(&workspace.db, file);
-
         let edits: Vec<TextEdit> = fix
             .edits()
             .iter()
-            .map(|edit| TextEdit {
+            .map(|edit| edit_to_text_edit(workspace, file, edit))
-                range: Range::from_text_range(
-                    edit.range(),
-                    &index,
-                    &source,
-                    workspace.position_encoding,
-                ),
-                new_text: edit.content().unwrap_or_default().to_string(),
-            })
             .collect();
 
         let title = self
@@ -781,7 +815,21 @@ impl Diagnostic {
             .map(ToString::to_string)
             .unwrap_or_else(|| format!("Fix {}", self.inner.id()));
 
-        Some(CodeAction { title, edits })
+        Some(CodeAction {
+            title,
+            edits,
+            preferred: true,
+        })
+    }
+}
+
+fn edit_to_text_edit(workspace: &Workspace, file: File, edit: &Edit) -> TextEdit {
+    let source = source_text(&workspace.db, file);
+    let index = line_index(&workspace.db, file);
+
+    TextEdit {
+        range: Range::from_text_range(edit.range(), &index, &source, workspace.position_encoding),
+        new_text: edit.content().unwrap_or_default().to_string(),
     }
 }
 
@@ -793,6 +841,7 @@ pub struct CodeAction {
     pub title: String,
     #[wasm_bindgen(getter_with_clone)]
     pub edits: Vec<TextEdit>,
+    pub preferred: bool,
 }
 
 #[wasm_bindgen]

playground/ty/src/Editor/Editor.tsx

@@ -607,6 +607,10 @@ class PlaygroundServer
     _token: CancellationToken,
   ): languages.ProviderResult<languages.CodeActionList> {
     const actions: languages.CodeAction[] = [];
+    const fileHandle = this.getFileHandleForModel(model);
+    if (fileHandle == null) {
+      return undefined;
+    }
 
     for (const diagnostic of this.diagnostics) {
       const diagnosticRange = diagnostic.range;
@@ -619,15 +623,19 @@ class PlaygroundServer
         continue;
       }
 
-      const codeAction = diagnostic.raw.codeAction(this.props.workspace);
+      const codeActions = this.props.workspace.codeActions(
-      if (codeAction == null) {
+        fileHandle,
+        diagnostic.raw,
+      );
+      if (codeActions == null) {
         continue;
       }
 
+      for (const codeAction of codeActions) {
         actions.push({
           title: codeAction.title,
           kind: "quickfix",
-        isPreferred: true,
+          isPreferred: codeAction.preferred,
           edit: {
             edits: codeAction.edits.map((edit) => ({
               resource: model.uri,
@@ -640,6 +648,7 @@ class PlaygroundServer
           },
         });
       }
+    }
 
     if (actions.length === 0) {
       return undefined;