Fixes env vars in languages other than Python

template/server/contexts.py

@@ -8,6 +8,7 @@
 from consts import JUPYTER_BASE_URL
 from errors import ExecutionError
 from messaging import ContextWebSocket
+from envs import get_envs
 
 logger = logging.Logger(__name__)
 
@@ -51,6 +52,7 @@ async def create_context(client, websockets: dict, language: str, cwd: str) -> C
     session_data = response.json()
     session_id = session_data["id"]
     context_id = session_data["kernel"]["id"]
+    global_env_vars = get_envs()
 
     logger.debug(f"Created context {context_id}")
 
@@ -67,4 +69,12 @@ async def create_context(client, websockets: dict, language: str, cwd: str) -> C
             status_code=500,
         )
 
+    try:
+        await ws.set_env_vars(global_env_vars)
+    except ExecutionError as e:
+        return PlainTextResponse(
+            "Failed to set environment variables",
+            status_code=500,
+        )
+
     return Context(language=language, id=context_id, cwd=cwd)

template/server/main.py

@@ -34,31 +34,22 @@ async def lifespan(app: FastAPI):
     global client
     client = httpx.AsyncClient()
 
-    with open("/root/.jupyter/kernel_id") as file:
-        default_context_id = file.read().strip()
-
-    default_ws = ContextWebSocket(
-        default_context_id,
-        str(uuid.uuid4()),
-        "python",
-        "/home/user",
-    )
-    default_websockets["python"] = default_context_id
-    websockets["default"] = default_ws
-    websockets[default_context_id] = default_ws
-
-    logger.info("Connecting to default runtime")
-    await default_ws.connect()
-
-    websockets["default"] = default_ws
+    try:
+        default_context = await create_context(client, websockets, "python", "/home/user")
+        default_websockets["python"] = default_context.id
+        websockets["default"] = websockets[default_context.id]
 
-    logger.info("Connected to default runtime")
-    yield
+        logger.info("Connected to default runtime")
+        yield
 
-    for ws in websockets.values():
-        await ws.close()
+        # Will cleanup after application shuts down
+        for ws in websockets.values():
+            await ws.close()
 
-    await client.aclose()
+        await client.aclose()
+    except Exception as e:
+        logger.error(f"Failed to initialize default context: {e}")
+        raise
 
 
 app = FastAPI(lifespan=lifespan)

template/server/messaging.py

@@ -135,6 +135,69 @@ async def _wait_for_result(self, message_id: str):
 
             yield output.model_dump(exclude_none=True)
 
+    async def set_env_vars(self, env_vars: Dict[StrictStr, str]):
+        env_commands = []
+
+        for k, v in env_vars.items():
+            if self.language == "python":
+                env_commands.append(f"os.environ['{k}'] = '{v}'")
+            elif self.language in ["javascript", "typescript"]:
+                env_commands.append(f"process.env['{k}'] = '{v}'")
+            elif self.language == "deno":
+                env_commands.append(f"Deno.env.set('{k}', '{v}')")
+            elif self.language == "r":
+                env_commands.append(f"Sys.setenv('{k}' = '{v}')")
+            elif self.language == "java":
+                env_commands.append(f"System.setProperty('{k}', '{v}')")
+            elif self.language == "bash":
+                env_commands.append(f"export {k}='{v}'")
+
+        if env_commands:
+            env_vars_snippet = "\n".join(env_commands)
+            print(f"Setting env vars: {env_vars_snippet}")
+            request = self._get_execute_request(str(uuid.uuid4()), env_vars_snippet, False)
+            await self._ws.send(request)
+
+    async def reset_env_vars(self, env_vars: Dict[StrictStr, str]):
+        global_env_vars = get_envs()
+
+        # Create a dict of vars to reset and a list of vars to remove
+        vars_to_reset = {}
+        vars_to_remove = []
+
+        for key in env_vars:
+            if key in global_env_vars:
+                vars_to_reset[key] = global_env_vars[key]
+            else:
+                vars_to_remove.append(key)
+
+        # Reset vars that exist in global env vars
+        if vars_to_reset:
+            await self.set_env_vars(vars_to_reset)
+
+        # Remove vars that don't exist in global env vars
+        if vars_to_remove:
+            remove_commands = []
+            for key in vars_to_remove:
+                if self.language == "python":
+                    remove_commands.append(f"del os.environ['{key}']")
+                elif self.language in ["javascript", "typescript"]:
+                    remove_commands.append(f"delete process.env['{key}']")
+                elif self.language == "deno":
+                    remove_commands.append(f"Deno.env.delete('{key}')")
+                elif self.language == "r":
+                    remove_commands.append(f"Sys.unsetenv('{key}')")
+                elif self.language == "java":
+                    remove_commands.append(f"System.clearProperty('{key}')")
+                elif self.language == "bash":
+                    remove_commands.append(f"unset {key}")
+
+            if remove_commands:
+                remove_snippet = "\n".join(remove_commands)
+                print(f"Removing env vars: {remove_snippet}")
+                request = self._get_execute_request(str(uuid.uuid4()), remove_snippet, False)
+                await self._ws.send(request)
+
     async def change_current_directory(
         self, path: Union[str, StrictStr], language: str
     ):
@@ -178,26 +241,10 @@ async def execute(
         if self._ws is None:
             raise Exception("WebSocket not connected")
 
-        global_env_vars = get_envs()
-        env_vars = {**global_env_vars, **env_vars} if env_vars else global_env_vars
         async with self._lock:
+            # set env vars (will override global env vars)
             if env_vars:
-                vars_to_set = {**global_env_vars, **env_vars}
-
-                # if there is an indent in the code, we need to add the env vars at the beginning of the code
-                lines = code.split("\n")
-                indent = 0
-                for i, line in enumerate(lines):
-                    if line.strip() != "":
-                        indent = len(line) - len(line.lstrip())
-                        break
-
-                if self.language == "python":
-                    code = (
-                        indent * " "
-                        + f"os.environ.set_envs_for_execution({vars_to_set})\n"
-                        + code
-                    )
+                await self.set_env_vars(env_vars)
 
             logger.info(code)
             request = self._get_execute_request(message_id, code, False)
@@ -211,6 +258,10 @@ async def execute(
 
             del self._executions[message_id]
 
+            # reset env vars to their previous values, if they were set globally or remove them
+            if env_vars:
+                await self.reset_env_vars(env_vars)
+
     async def _receive_message(self):
         if not self._ws:
             logger.error("No WebSocket connection")
@@ -276,15 +327,15 @@ async def _process_message(self, data: dict):
 
         elif data["msg_type"] == "stream":
             if data["content"]["name"] == "stdout":
-                logger.debug(f"Execution {parent_msg_ig} received stdout")
+                logger.debug(f"Execution {parent_msg_ig} received stdout: {data['content']['text']}")
                 await queue.put(
                     Stdout(
                         text=data["content"]["text"], timestamp=data["header"]["date"]
                     )
                 )
 
             elif data["content"]["name"] == "stderr":
-                logger.debug(f"Execution {parent_msg_ig} received stderr")
+                logger.debug(f"Execution {parent_msg_ig} received stderr: {data['content']['text']}")
                 await queue.put(
                     Stderr(
                         text=data["content"]["text"], timestamp=data["header"]["date"]

template/start-up.sh

@@ -13,18 +13,6 @@ function start_jupyter_server() {
 		response=$(curl -s -o /dev/null -w "%{http_code}" "http://localhost:8888/api/status")
 	done
 
-	response=$(curl -s -X POST "localhost:8888/api/sessions" -H "Content-Type: application/json" -d '{"path": "/home/user", "kernel": {"name": "python3"}, "type": "notebook", "name": "default"}')
-	status=$(echo "${response}" | jq -r '.kernel.execution_state')
-	if [[ ${status} != "starting" ]]; then
-		echo "Error creating kernel: ${response} ${status}"
-		exit 1
-	fi
-
-	sudo mkdir -p /root/.jupyter
-	kernel_id=$(echo "${response}" | jq -r '.kernel.id')
-	sudo echo "${kernel_id}" | sudo tee /root/.jupyter/kernel_id >/dev/null
-	sudo echo "${response}" | sudo tee /root/.jupyter/.session_info >/dev/null
-
 	cd /root/.server/
 	/root/.server/.venv/bin/uvicorn main:app --host 0.0.0.0 --port 49999 --workers 1 --no-access-log --no-use-colors --timeout-keep-alive 640
 }