From 18f63aab40d429c3b5ad3d495da5144f085da94a Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Sat, 6 Jul 2024 16:15:06 +0200
Subject: [PATCH] Add environment to runnable lsp extension

---
 crates/rust-analyzer/src/handlers/request.rs | 10 +++---
 crates/rust-analyzer/src/lsp/ext.rs          | 11 +++---
 crates/rust-analyzer/src/lsp/to_proto.rs     | 11 +++---
 docs/dev/lsp-extensions.md                   | 37 +++++++++++++++++---
 editors/code/src/lsp_ext.ts                  | 34 +++++++++++++++---
 editors/code/src/run.ts                      | 14 ++++----
 6 files changed, 84 insertions(+), 33 deletions(-)

diff --git a/crates/rust-analyzer/src/handlers/request.rs b/crates/rust-analyzer/src/handlers/request.rs
index 8e39b15da3dd..7e0db8edd818 100644
--- a/crates/rust-analyzer/src/handlers/request.rs
+++ b/crates/rust-analyzer/src/handlers/request.rs
@@ -847,7 +847,7 @@ pub(crate) fn handle_runnables(
             if expect_test {
                 if let lsp_ext::RunnableArgs::Cargo(r) = &mut runnable.args {
                     runnable.label = format!("{} + expect", runnable.label);
-                    r.expect_test = Some(true);
+                    r.environment.insert("UPDATE_EXPECT".to_owned(), "1".to_owned());
                 }
             }
             res.push(runnable);
@@ -884,12 +884,12 @@ pub(crate) fn handle_runnables(
                     kind: lsp_ext::RunnableKind::Cargo,
                     args: lsp_ext::RunnableArgs::Cargo(lsp_ext::CargoRunnableArgs {
                         workspace_root: Some(spec.workspace_root.clone().into()),
-                        cwd: Some(cwd.into()),
+                        cwd: cwd.into(),
                         override_cargo: config.override_cargo.clone(),
                         cargo_args,
                         cargo_extra_args: config.cargo_extra_args.clone(),
                         executable_args: Vec::new(),
-                        expect_test: None,
+                        environment: Default::default(),
                     }),
                 })
             }
@@ -903,12 +903,12 @@ pub(crate) fn handle_runnables(
                     kind: lsp_ext::RunnableKind::Cargo,
                     args: lsp_ext::RunnableArgs::Cargo(lsp_ext::CargoRunnableArgs {
                         workspace_root: None,
-                        cwd: None,
+                        cwd: ".".into(),
                         override_cargo: config.override_cargo,
                         cargo_args: vec!["check".to_owned(), "--workspace".to_owned()],
                         cargo_extra_args: config.cargo_extra_args,
                         executable_args: Vec::new(),
-                        expect_test: None,
+                        environment: Default::default(),
                     }),
                 });
             }
diff --git a/crates/rust-analyzer/src/lsp/ext.rs b/crates/rust-analyzer/src/lsp/ext.rs
index b82ba4419041..3766e3f46a7d 100644
--- a/crates/rust-analyzer/src/lsp/ext.rs
+++ b/crates/rust-analyzer/src/lsp/ext.rs
@@ -460,28 +460,27 @@ pub enum RunnableKind {
 #[derive(Deserialize, Serialize, Debug)]
 #[serde(rename_all = "camelCase")]
 pub struct CargoRunnableArgs {
-    // command to be executed instead of cargo
+    pub environment: FxHashMap<String, String>,
+    pub cwd: Utf8PathBuf,
+    /// Command to be executed instead of cargo
     pub override_cargo: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub workspace_root: Option<Utf8PathBuf>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub cwd: Option<Utf8PathBuf>,
     // command, --package and --lib stuff
     pub cargo_args: Vec<String>,
     // user-specified additional cargo args, like `--release`.
     pub cargo_extra_args: Vec<String>,
     // stuff after --
     pub executable_args: Vec<String>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub expect_test: Option<bool>,
 }
 
 #[derive(Deserialize, Serialize, Debug)]
 #[serde(rename_all = "camelCase")]
 pub struct ShellRunnableArgs {
+    pub environment: FxHashMap<String, String>,
+    pub cwd: Utf8PathBuf,
     pub program: String,
     pub args: Vec<String>,
-    pub cwd: Utf8PathBuf,
 }
 
 pub enum RelatedTests {}
diff --git a/crates/rust-analyzer/src/lsp/to_proto.rs b/crates/rust-analyzer/src/lsp/to_proto.rs
index db5f666a5b95..fe65d1fc4fdd 100644
--- a/crates/rust-analyzer/src/lsp/to_proto.rs
+++ b/crates/rust-analyzer/src/lsp/to_proto.rs
@@ -15,7 +15,7 @@ use ide::{
 };
 use ide_db::{rust_doc::format_docs, FxHasher};
 use itertools::Itertools;
-use paths::{Utf8Component, Utf8Prefix};
+use paths::{Utf8Component, Utf8PathBuf, Utf8Prefix};
 use semver::VersionReq;
 use serde_json::to_value;
 use vfs::AbsPath;
@@ -1390,10 +1390,10 @@ pub(crate) fn runnable(
                     workspace_root: Some(workspace_root.into()),
                     override_cargo: config.override_cargo,
                     cargo_args,
-                    cwd: Some(cwd.into()),
+                    cwd: cwd.into(),
                     cargo_extra_args: config.cargo_extra_args,
                     executable_args,
-                    expect_test: None,
+                    environment: Default::default(),
                 }),
             }))
         }
@@ -1407,6 +1407,7 @@ pub(crate) fn runnable(
                         program: json_shell_runnable_args.program,
                         args: json_shell_runnable_args.args,
                         cwd: json_shell_runnable_args.cwd,
+                        environment: Default::default(),
                     };
                     Ok(Some(lsp_ext::Runnable {
                         label,
@@ -1433,10 +1434,10 @@ pub(crate) fn runnable(
                     workspace_root: None,
                     override_cargo: config.override_cargo,
                     cargo_args,
-                    cwd: None,
+                    cwd: Utf8PathBuf::from("."),
                     cargo_extra_args: config.cargo_extra_args,
                     executable_args,
-                    expect_test: None,
+                    environment: Default::default(),
                 }),
             }))
         }
diff --git a/docs/dev/lsp-extensions.md b/docs/dev/lsp-extensions.md
index 695fec7e8e01..1cafcae7bb61 100644
--- a/docs/dev/lsp-extensions.md
+++ b/docs/dev/lsp-extensions.md
@@ -1,5 +1,5 @@
 <!---
-lsp/ext.rs hash: 8e6e340f2899b5e9
+lsp/ext.rs hash: 3605fab9e66e14a0
 
 If you need to change the above hash to make the test pass, please check if you
 need to adjust this doc as well and ping this issue:
@@ -376,12 +376,34 @@ rust-analyzer supports two `kind`s of runnables, `"cargo"` and `"shell"`. The `a
 
 ```typescript
 {
+    /**
+     * Environment variables to set before running the command.
+     */
+    environment: Record<string, string>;
+    /**
+     * The working directory to run the command in.
+     */
+    cwd: string;
+    /**
+     * The workspace root directory of the cargo project.
+     */
     workspaceRoot?: string;
-    cwd?: string;
+    /**
+     * The cargo command to run.
+     */
     cargoArgs: string[];
+    /**
+     * Extra arguments to pass to cargo.
+     */
+    // What is the point of this when cargoArgs exists?
     cargoExtraArgs: string[];
+    /**
+     * Arguments to pass to the executable, these will be passed to the command after a `--` argument.
+     */
     executableArgs: string[];
-    expectTest?: boolean;
+    /**
+     * Command to execute instead of `cargo`.
+     */
     overrideCargo?: string;
 }
 ```
@@ -390,10 +412,17 @@ The args for `"shell"` look like this:
 
 ```typescript
 {
+    /**
+     * Environment variables to set before running the command.
+     */
+    environment: Record<string, string>;
+    /**
+     * The working directory to run the command in.
+     */
+    cwd: string;
     kind: string;
     program: string;
     args: string[];
-    cwd: string;
 }
 ```
 
diff --git a/editors/code/src/lsp_ext.ts b/editors/code/src/lsp_ext.ts
index 699052e4d447..fee023bbc94a 100644
--- a/editors/code/src/lsp_ext.ts
+++ b/editors/code/src/lsp_ext.ts
@@ -235,22 +235,46 @@ type RunnableShell = {
     args: ShellRunnableArgs;
 };
 
+export type CommonRunnableArgs = {
+    /**
+     * Environment variables to set before running the command.
+     */
+    environment: Record<string, string>;
+    /**
+     * The working directory to run the command in.
+     */
+    cwd: string;
+};
+
 export type ShellRunnableArgs = {
     kind: string;
     program: string;
     args: string[];
-    cwd: string;
-};
+} & CommonRunnableArgs;
 
 export type CargoRunnableArgs = {
+    /**
+     * The workspace root directory of the cargo project.
+     */
     workspaceRoot?: string;
+    /**
+     * The cargo command to run.
+     */
     cargoArgs: string[];
-    cwd: string;
+    /**
+     * Extra arguments to pass to cargo.
+     */
+    // What is the point of this when cargoArgs exists?
     cargoExtraArgs: string[];
+    /**
+     * Arguments to pass to the executable, these will be passed to the command after a `--` argument.
+     */
     executableArgs: string[];
-    expectTest?: boolean;
+    /**
+     * Command to execute instead of `cargo`.
+     */
     overrideCargo?: string;
-};
+} & CommonRunnableArgs;
 
 export type RunnablesParams = {
     textDocument: lc.TextDocumentIdentifier;
diff --git a/editors/code/src/run.ts b/editors/code/src/run.ts
index 7a9049af0ded..97dab8d9f65d 100644
--- a/editors/code/src/run.ts
+++ b/editors/code/src/run.ts
@@ -66,9 +66,12 @@ export class RunnableQuickPick implements vscode.QuickPickItem {
     }
 }
 
-export function prepareBaseEnv(): Record<string, string> {
+export function prepareBaseEnv(base?: Record<string, string>): Record<string, string> {
     const env: Record<string, string> = { RUST_BACKTRACE: "short" };
-    Object.assign(env, process.env as { [key: string]: string });
+    Object.assign(env, process.env);
+    if (base) {
+        Object.assign(env, base);
+    }
     return env;
 }
 
@@ -77,12 +80,7 @@ export function prepareEnv(
     runnableArgs: ra.CargoRunnableArgs,
     runnableEnvCfg: RunnableEnvCfg,
 ): Record<string, string> {
-    const env = prepareBaseEnv();
-
-    if (runnableArgs.expectTest) {
-        env["UPDATE_EXPECT"] = "1";
-    }
-
+    const env = prepareBaseEnv(runnableArgs.environment);
     const platform = process.platform;
 
     const checkPlatform = (it: RunnableEnvCfgItem) => {