[vscode] 7/n listeners and .env

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
@@ -381,6 +394,12 @@ export async function setupEnvironmentVariables(
       "Please define your environment variables."
     );
   }
+
+  // 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());
 }
 
 /**