use server info

Author: asukaminato0721

lsp/src/codeLens.ts

@@ -19,10 +19,12 @@ type CodeLensPosition = {
 
 type RunMainArgs = {
   uri?: string;
+  cwd?: string;
 };
 
 type RunTestArgs = {
   uri?: string;
+  cwd?: string;
   position?: CodeLensPosition;
   testName?: string;
   className?: string;
@@ -45,10 +47,6 @@ function scopeForUri(uri: vscode.Uri): vscode.WorkspaceFolder | vscode.TaskScope
   return vscode.workspace.getWorkspaceFolder(uri) ?? vscode.TaskScope.Workspace;
 }
 
-function cwdForUri(uri: vscode.Uri): string | undefined {
-  return vscode.workspace.getWorkspaceFolder(uri)?.uri.fsPath;
-}
-
 function moduleNameFromPath(uri: vscode.Uri): string | undefined {
   const workspaceFolder = vscode.workspace.getWorkspaceFolder(uri);
   if (!workspaceFolder) {
@@ -85,6 +83,7 @@ async function runAtCursor(
 
 async function executeProcessTask(
   uri: vscode.Uri,
+  cwd: string | undefined,
   definition: vscode.TaskDefinition,
   label: string,
   command: string,
@@ -98,7 +97,7 @@ async function executeProcessTask(
     label,
     TASK_SOURCE,
     new vscode.ProcessExecution(command, args, {
-      cwd: cwdForUri(uri),
+      cwd,
     }),
   );
   task.presentationOptions = {
@@ -128,6 +127,7 @@ async function runMainFile(
   }
   await executeProcessTask(
     uri,
+    args.cwd,
     {type: TASK_SOURCE, action: 'runMain'},
     'Pyrefly: Run File',
     interpreter,
@@ -143,6 +143,7 @@ async function runTestAtLocation(
     return;
   }
   const uri = vscode.Uri.parse(args.uri);
+  const cwd = args.cwd;
   const className = args.className;
   const testName = args.testName;
 
@@ -175,6 +176,7 @@ async function runTestAtLocation(
     }
     await executeProcessTask(
       uri,
+      cwd,
       {type: TASK_SOURCE, action: 'runUnittest'},
       'Pyrefly: Run Test',
       interpreter,
@@ -192,6 +194,7 @@ async function runTestAtLocation(
   }
   await executeProcessTask(
     uri,
+    cwd,
     {type: TASK_SOURCE, action: 'runPytest'},
     'Pyrefly: Run Test',
     interpreter,

pyrefly/lib/lsp/non_wasm/server.rs

@@ -2405,6 +2405,24 @@ impl Server {
         telemetry.set_file_stats(TelemetryFileStats { uri, config_root });
     }
 
+    fn runnable_code_lens_cwd(&self, path: &std::path::Path) -> Option<String> {
+        let config = self.state.config_finder().python_file(
+            ModuleNameWithKind::guaranteed(ModuleName::unknown()),
+            &ModulePath::filesystem(path.to_path_buf()),
+        );
+        let cwd = config
+            .source
+            .root()
+            .map(std::path::Path::to_path_buf)
+            .or_else(|| {
+                self.workspaces
+                    .get_with(path.to_path_buf(), |(workspace_root, _)| {
+                        workspace_root.cloned()
+                    })
+            })?;
+        Some(cwd.to_string_lossy().into_owned())
+    }
+
     fn send_response(&self, x: Response) {
         self.connection.send(Message::Response(x))
     }
@@ -4681,14 +4699,15 @@ impl Server {
         let path = self.path_for_uri(uri)?;
         let runnable_code_lens = self
             .workspaces
-            .get_with(path, |(_, workspace)| workspace.runnable_code_lens);
+            .get_with(path.clone(), |(_, workspace)| workspace.runnable_code_lens);
         if !runnable_code_lens {
             return Some(Vec::new());
         }
         let maybe_cell_idx = self.maybe_get_cell_index(uri);
         let handle = self.make_handle_if_enabled(uri, Some(CodeLensRequest::METHOD))?;
         let info = transaction.get_module_info(&handle)?;
         let entries = transaction.code_lens_entries(&handle)?;
+        let cwd = self.runnable_code_lens_cwd(&path);
 
         let mut lenses = Vec::new();
         for entry in entries {
@@ -4697,14 +4716,20 @@ impl Server {
             }
             let range = info.to_lsp_range(entry.range);
             let (title, command, arguments) = match entry.kind {
-                CodeLensKind::Run => (
-                    "Run",
-                    "pyrefly.runMain",
-                    Some(vec![serde_json::json!({ "uri": uri.to_string() })]),
-                ),
+                CodeLensKind::Run => {
+                    let mut args = serde_json::Map::new();
+                    args.insert("uri".to_owned(), serde_json::json!(uri.to_string()));
+                    if let Some(cwd) = &cwd {
+                        args.insert("cwd".to_owned(), serde_json::json!(cwd));
+                    }
+                    ("Run", "pyrefly.runMain", Some(vec![Value::Object(args)]))
+                }
                 CodeLensKind::Test => {
                     let mut args = serde_json::Map::new();
                     args.insert("uri".to_owned(), serde_json::json!(uri.to_string()));
+                    if let Some(cwd) = &cwd {
+                        args.insert("cwd".to_owned(), serde_json::json!(cwd));
+                    }
                     args.insert(
                         "position".to_owned(),
                         serde_json::json!({

pyrefly/lib/test/lsp/lsp_interaction/code_lens.rs

@@ -65,7 +65,7 @@ fn test_code_lens_for_tests_and_main() {
                     has_run |= lens.range.start.line == 17;
                 }
                 if command.command == "pyrefly.runTest" && command.title == "Test" {
-                    let args = command.arguments.unwrap_or_default();
+                    let args = command.arguments.clone().unwrap_or_default();
                     let Some(Value::Object(obj)) = args.get(0) else {
                         continue;
                     };
@@ -91,6 +91,56 @@ fn test_code_lens_for_tests_and_main() {
     interaction.shutdown().unwrap();
 }
 
+#[test]
+fn test_code_lens_uses_config_root_for_cwd() {
+    let root = get_test_files_root();
+    let mut interaction = LspInteraction::new();
+    let test_root = root.path().join("code_lens");
+    interaction.set_root(test_root.clone());
+    interaction
+        .initialize(InitializeSettings {
+            configuration: Some(Some(runnable_code_lens_config())),
+            ..Default::default()
+        })
+        .unwrap();
+
+    interaction
+        .client
+        .did_open("nested_project/main_and_tests.py");
+
+    let path = test_root.join("nested_project/main_and_tests.py");
+    let uri = Url::from_file_path(&path).unwrap();
+    let expected_cwd = test_root
+        .join("nested_project")
+        .to_string_lossy()
+        .into_owned();
+
+    interaction
+        .client
+        .send_request::<CodeLensRequest>(json!({
+            "textDocument": {
+                "uri": uri.to_string()
+            },
+        }))
+        .expect_response_with(|response: Option<Vec<CodeLens>>| {
+            let Some(lenses) = response else {
+                return false;
+            };
+            let mut saw_lens = false;
+            lenses.into_iter().all(|lens| {
+                saw_lens = true;
+                lens.command
+                    .and_then(|command| command.arguments)
+                    .and_then(|args| args.into_iter().next())
+                    .and_then(|arg| arg.get("cwd").and_then(Value::as_str).map(str::to_owned))
+                    .is_some_and(|cwd| cwd == expected_cwd)
+            }) && saw_lens
+        })
+        .unwrap();
+
+    interaction.shutdown().unwrap();
+}
+
 #[test]
 fn test_code_lens_ignores_stub_files() {
     let root = get_test_files_root();

pyrefly/lib/test/lsp/lsp_interaction/test_files/code_lens/nested_project/main_and_tests.py

@@ -0,0 +1,5 @@
+def test_nested():
+    pass
+
+if __name__ == "__main__":
+    test_nested()

pyrefly/lib/test/lsp/lsp_interaction/test_files/code_lens/nested_project/pyrefly.toml

@@ -0,0 +1 @@
+[tool.pyrefly]