[vscode] 8/n Remove LCA env checks (#1398)

[vscode] 8/n Remove LCA env checks

Summary:

Now that a user can specify a .env file path with no restrictions on its
location, the LCA checks need to be removed.

LCA = Lowest Common Ancestor



Test Plan:

---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with
[ReviewStack](https://reviewstack.dev/lastmile-ai/aiconfig/pull/1398).
* __->__ #1398
* #1390
* #1387
* #1389

python/src/aiconfig/editor/server/server_utils.py

@@ -10,7 +10,7 @@
 from textwrap import dedent
 from threading import Event
 from types import ModuleType
-from typing import Any, Callable, NewType, Type, TypeVar, cast
+from typing import Any, Callable, NewType, Type, TypeVar, cast, Optional
 
 import lastmile_utils.lib.core.api as core_utils
 import result
@@ -66,6 +66,7 @@ class EditServerConfig(core_utils.Record):
     log_level: str | int = "INFO"
     server_mode: ServerMode = ServerMode.PROD
     parsers_module_path: str = "aiconfig_model_registry.py"
+    env_file_path: Optional[str] = None
 
     @field_validator("server_mode", mode="before")
     def convert_to_mode(
@@ -84,6 +85,7 @@ class StartServerConfig(core_utils.Record):
     log_level: str | int = "INFO"
     server_mode: ServerMode = ServerMode.PROD
     parsers_module_path: str = "aiconfig_model_registry.py"
+    env_file_path: Optional[str] = None
 
     @field_validator("server_mode", mode="before")
     def convert_to_mode(
@@ -294,13 +296,12 @@ def init_server_state(
     aiconfigrc_path: str,
 ) -> Result[None, str]:
     LOGGER.info("Initializing server state for 'edit' command")
-    # TODO: saqadri - load specific .env file if specified
-    dotenv.load_dotenv()
     _load_user_parser_module_if_exists(
         initialization_settings.parsers_module_path
     )
     state = get_server_state(app)
     state.aiconfigrc_path = aiconfigrc_path
+    state.env_file_path = initialization_settings.env_file_path
 
     if isinstance(initialization_settings, StartServerConfig):
         # The aiconfig will be loaded later, when the editor sends the payload.

vscode-extension/package.json

@@ -128,6 +128,11 @@
           "type": "string",
           "default": "",
           "description": "Version of the AIConfig Editor extension last time it was activated. We use this value to check if the extension has been updated to prompt users to reload VS Code"
+        },
+        "vscode-aiconfig.env_file_path": {
+          "type": "string",
+          "default": "",
+          "description": "File path to your .env file containing your API keys"
         }
       }
     },

vscode-extension/src/constants.ts

@@ -3,3 +3,5 @@ export const PYTHON_INTERPRETER_CACHE_KEY_NAME = "pythonInterpreter";
 // Used for prompting user to reload VS Code window on update, and ensuring
 // that walkthrough gets shown on first install
 export const VERSION_KEY_NAME = "version";
+
+export const ENV_FILE_PATH = "env_file_path";

vscode-extension/src/editor_server/editorServer.ts

@@ -3,6 +3,7 @@ import { EXTENSION_NAME } from "../util";
 import { getPythonPath } from "../utilities/pythonSetupUtils";
 import { ChildProcessWithoutNullStreams, spawn } from "child_process";
 import serverPortManager from "./serverPortManager";
+import { ENV_FILE_PATH } from "../constants";
 
 export enum EditorServerState {
   Starting = "Starting",
@@ -73,6 +74,8 @@ export class EditorServer {
       const modelRegistryPathArgs = modelRegistryPath
         ? ["--parsers-module-path", modelRegistryPath]
         : [];
+      const filePath = config.get(ENV_FILE_PATH) as string;
+      const envFilePathArgs = filePath ? ["--env-file-path", filePath] : [];
 
       this.port = await serverPortManager.getPort();
 
@@ -89,6 +92,7 @@ export class EditorServer {
           "--server-port",
           this.port.toString(),
           ...modelRegistryPathArgs,
+          ...envFilePathArgs,
         ],
         {
           cwd: this.cwd,

vscode-extension/src/extension.ts

@@ -27,13 +27,15 @@ import {
   setupEnvironmentVariables,
   getConfigurationTarget,
   getModeFromDocument,
+  updateServerEnv
 } from "./util";
 import {
   initialize,
   installDependencies,
   savePythonInterpreterToCache,
 } from "./utilities/pythonSetupUtils";
 import { performVersionInstallAndUpdateActionsIfNeeded } from "./utilities/versionUpdateUtils";
+import { ENV_FILE_PATH } from "./constants";
 
 // This method is called when your extension is activated
 // Your extension is activated the very first time a command is executed
@@ -208,6 +210,25 @@ export async function activate(context: vscode.ExtensionContext) {
       }
     )
   );
+
+  // Handle changes to the .env path
+  context.subscriptions.push(vscode.workspace.onDidChangeConfiguration( async (event) => {
+    if (event.affectsConfiguration("vscode-aiconfig" + "." + ENV_FILE_PATH)) {
+      // Get new env
+      const envPath = vscode.workspace.getConfiguration("vscode-aiconfig").get(ENV_FILE_PATH) as string;
+      console.log(`New .env path set: ${envPath}`);
+      // set env on all AIConfig Servers
+      const editors: Array<AIConfigEditorState> = Array.from(
+        aiconfigEditorManager.getRegisteredEditors()
+      );
+      if (editors.length > 0) {
+        editors.forEach(async (editor) => {
+          await updateServerEnv(editor.editorServer.url, envPath);
+        });
+
+      }
+  }
+  }));
 }
 
 // This method is called when your extension is deactivated

vscode-extension/src/util.ts

@@ -8,6 +8,8 @@ import fs from "fs";
 import path from "path";
 import os from "os";
 
+import { ENV_FILE_PATH } from "./constants";
+
 export const EXTENSION_NAME = "vscode-aiconfig";
 export const COMMANDS = {
   INIT: `${EXTENSION_NAME}.init`,
@@ -45,6 +47,8 @@ export const EDITOR_SERVER_ROUTE_TABLE = {
     urlJoin(hostUrl, EDITOR_SERVER_API_ENDPOINT, "/load_content"),
   LOAD_MODEL_PARSER_MODULE: (hostUrl: string) =>
     urlJoin(hostUrl, EDITOR_SERVER_API_ENDPOINT, "/load_model_parser_module"),
+  SET_ENV_FILE_PATH: (hostUrl: string) =>
+    urlJoin(hostUrl, EDITOR_SERVER_API_ENDPOINT, "/set_env_file_path"),
 };
 
 export async function isServerReady(serverUrl: string) {
@@ -98,6 +102,15 @@ export async function updateServerState(
   });
 }
 
+export async function updateServerEnv(serverUrl: string, filePath: string) {
+  return await ufetch.post(
+    EDITOR_SERVER_ROUTE_TABLE.SET_ENV_FILE_PATH(serverUrl),
+    {
+      [ENV_FILE_PATH]: filePath,
+    }
+  );
+}
+
 // Figure out what kind of AIConfig this is that we are loading
 export function getModeFromDocument(
   document: vscode.TextDocument
@@ -317,12 +330,11 @@ export async function setupEnvironmentVariables(
 ) {
   const homedir = require("os").homedir(); // This is cross-platform: https://stackoverflow.com/a/9081436
   const defaultEnvPath = path.join(homedir, ".env");
-  const lowestCommonWorkspacePath = getLowestCommonAncestorAcrossWorkspaces();
 
   const envPath = await vscode.window.showInputBox({
     prompt: "Enter the path of your .env file",
     value: defaultEnvPath,
-    validateInput: (input) => validateEnvPath(input, lowestCommonWorkspacePath),
+    validateInput: (input) => validateEnvPath(input),
   });
 
   if (!envPath) {
@@ -381,71 +393,21 @@ export async function setupEnvironmentVariables(
       "Please define your environment variables."
     );
   }
-}
-
-/**
- * Some VS Code setups can have multiple workspaces, in which
- * case we should take the lowest common ancestor path that is shared
- * across all of them so that the same .env file can be used for multiple
- * AIConfig files
- * @returns lowestCommonAncestorPath (string | undefined)
- *    -> string of path to lowest common ancestor: empty means no shared path
- *    -> undefined if no workspaces are defined in VS Code session
- */
-function getLowestCommonAncestorAcrossWorkspaces(): string | undefined {
-  const workspaceFolders = vscode.workspace.workspaceFolders;
-  if (workspaceFolders === undefined || workspaceFolders.length === 0) {
-    return undefined;
-  }
 
-  const workspacePaths = workspaceFolders.map((folder) =>
-    path.normalize(folder.uri.fsPath)
-  );
-  let lowestCommonAncestorPath: string;
-  const separator = path.sep; // Handles Windows and Linux
-  lowestCommonAncestorPath = workspacePaths.reduce(
-    (currLowestCommonAncestorPath, currPath) => {
-      const ancestorFolders = currLowestCommonAncestorPath.split(separator);
-      const currPathFolders = currPath.split(separator);
-      const commonPathFolders: Array<string> = [];
-      for (var i = 0; i < ancestorFolders.length; i++) {
-        if (ancestorFolders[i] === currPathFolders[i]) {
-          commonPathFolders.push(ancestorFolders[i]);
-        } else {
-          break;
-        }
-      }
-      return commonPathFolders.join(separator);
-    }
-  );
-  return lowestCommonAncestorPath;
+  // Update Server Env FLow
+  // Set the .env file path in the settings
+  // vscode Extension has a listener for changes defined at activation.
+  const config = vscode.workspace.getConfiguration(EXTENSION_NAME);
+  await config.update(ENV_FILE_PATH, envPath, getConfigurationTarget());
 }
 
 function validateEnvPath(
   inputPath: string,
-  workspacePath: string | undefined
 ): string | null {
   if (!inputPath) {
     return "File path is required";
   } else if (path.basename(inputPath) !== ".env") {
     return 'Filename must be ".env"';
-  } else if (workspacePath != null && workspacePath !== "") {
-    // loadenv() from Python checks each folder from the file/program where
-    // it's invoked for the presence of an `.env` file. Therefore, the `.env
-    // file must be saved either at the top-level directory of the workspace
-    // directory, or one of it's parent directories. This will ensure that if
-    // two AIConfig files are contained in separate paths within the workspace
-    // they'll still be able to access the same `.env` file.
-
-    // Note: If the `inputPath` directory is equal to the `workspacePath`,
-    // `relativePathFromEnvToWorkspace` will be an empty string
-    const relativePathFromEnvToWorkspace = path.relative(
-      path.dirname(inputPath),
-      workspacePath
-    );
-    if (relativePathFromEnvToWorkspace.startsWith("..")) {
-      return `File path must either be contained within the VS Code workspace directory ('${workspacePath}') or within a one of it's parent folders`;
-    }
-  }
+  } 
   return null;
 }