Guard vscode-specific logic with custom initialization options

src/server/mod.rs

@@ -47,9 +47,15 @@ mod imports;
 mod providers;
 mod symbols;
 
+#[derive(Default, serde::Serialize, serde::Deserialize)]
+pub struct InitializationOptions {
+    pub vscode_extension: bool,
+}
+
 struct ServerContext {
     pub storage: Arc<Mutex<DocumentStorage>>,
     pub analyzer: OnceLock<Arc<Analyzer>>,
+    pub options: OnceLock<Arc<InitializationOptions>>,
     pub client: TestableClient,
 }
 
@@ -58,6 +64,7 @@ impl ServerContext {
         Self {
             storage,
             analyzer: OnceLock::new(),
+            options: OnceLock::new(),
             client,
         }
     }
@@ -73,9 +80,13 @@ impl ServerContext {
             WorkspaceFinder::new(client_root),
             IndexingLevel::Disabled,
         )));
+        let options = OnceLock::new();
+        let _ = options.set(Default::default());
+
         Self {
             storage,
             analyzer,
+            options,
             client: TestableClient::new_for_testing(),
         }
     }
@@ -84,6 +95,7 @@ impl ServerContext {
         RequestContext {
             storage: self.storage.clone(),
             analyzer: self.analyzer.get().unwrap().clone(),
+            options: self.options.get().unwrap().clone(),
             client: self.client.clone(),
             request_time: Instant::now(),
         }
@@ -94,6 +106,7 @@ impl ServerContext {
 pub struct RequestContext {
     pub storage: Arc<Mutex<DocumentStorage>>,
     pub analyzer: Arc<Analyzer>,
+    pub options: Arc<InitializationOptions>,
     pub client: TestableClient,
     pub request_time: Instant,
 }
@@ -134,6 +147,12 @@ impl LanguageServer for Backend {
         ));
         self.context.analyzer.set(analyzer).ok();
 
+        let options: InitializationOptions = params
+            .initialization_options
+            .and_then(|value| serde_json::from_value(value).ok())
+            .unwrap_or_default();
+        self.context.options.set(Arc::new(options)).ok();
+
         Ok(InitializeResult {
             capabilities: ServerCapabilities {
                 text_document_sync: Some(TextDocumentSyncCapability::Kind(

src/server/providers/code_action.rs

@@ -67,22 +67,29 @@ async fn compute_import_actions(
     if imports.is_empty() {
         return Vec::new();
     }
-    if let Ok(only_import) = imports.iter().exactly_one() {
-        return vec![CodeActionOrCommand::CodeAction(CodeAction {
-            title: format!("Import `{name}` from `{only_import}`"),
-            kind: Some(CodeActionKind::QUICKFIX),
-            diagnostics: Some(vec![diagnostic.clone()]),
-            edit: Some(WorkspaceEdit {
-                changes: Some(HashMap::from([(
-                    Url::from_file_path(&current_file.document.path).unwrap(),
-                    vec![create_import_edit(&current_file, only_import)],
-                )])),
-                ..Default::default()
-            }),
-            command: None,
-            is_preferred: Some(true),
-            ..Default::default()
-        })];
+
+    let use_picker = context.options.vscode_extension && imports.len() >= 2;
+    if !use_picker {
+        return imports
+            .iter()
+            .map(|import| {
+                CodeActionOrCommand::CodeAction(CodeAction {
+                    title: format!("Import `{name}` from `{import}`"),
+                    kind: Some(CodeActionKind::QUICKFIX),
+                    diagnostics: Some(vec![diagnostic.clone()]),
+                    edit: Some(WorkspaceEdit {
+                        changes: Some(HashMap::from([(
+                            Url::from_file_path(&current_file.document.path).unwrap(),
+                            vec![create_import_edit(&current_file, import)],
+                        )])),
+                        ..Default::default()
+                    }),
+                    command: None,
+                    is_preferred: Some(true),
+                    ..Default::default()
+                })
+            })
+            .collect();
     }
 
     let data = ChooseImportCandidatesData {

src/server/providers/code_lens.rs

@@ -73,6 +73,9 @@ pub async fn code_lens(
     context: &RequestContext,
     params: CodeLensParams,
 ) -> Result<Option<Vec<CodeLens>>> {
+    if !context.options.vscode_extension {
+        return Ok(None);
+    }
     let configs = context.client.configurations().await;
     if !configs.target_lens {
         return Ok(None);

vscode-gn/src/extension.ts

@@ -31,6 +31,10 @@ import {
 
 const EXECUTABLE_SUFFIX: string = process.platform === 'win32' ? '.exe' : '';
 
+interface InitializationOptions {
+  vscode_extension: boolean;
+}
+
 function reportAsyncError(
   output: vscode.OutputChannel,
   result: Promise<void>
@@ -181,6 +185,9 @@ async function chooseImportCandidates(
 class GnLanguageClient extends LanguageClient {
   constructor(context: vscode.ExtensionContext, output: vscode.OutputChannel) {
     const clientOptions: LanguageClientOptions = {
+      initializationOptions: {
+        vscode_extension: true,
+      } as InitializationOptions,
       documentSelector: [
         {scheme: 'file', pattern: '**/*.gn'},
         {scheme: 'file', pattern: '**/*.gni'},