From 4fed7096f77142e75c14031dee6af05d2eba8b53 Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Wed, 20 Aug 2025 15:25:36 +0200
Subject: [PATCH 01/12] inline env var snippet

---
 template/server/messaging.py | 161 ++++++++++++++++-------------------
 1 file changed, 74 insertions(+), 87 deletions(-)

diff --git a/template/server/messaging.py b/template/server/messaging.py
index 7bc2902f..c6f6f222 100644
--- a/template/server/messaging.py
+++ b/template/server/messaging.py
@@ -114,6 +114,65 @@ def _get_execute_request(
             }
         )
 
+    def _get_env_var_command(self, key: str, value: str, action: str = "set") -> str:
+        """Get environment variable command for the current language."""
+        if action == "set":
+            if self.language == "python":
+                return f"import os; os.environ['{key}'] = '{value}'"
+            elif self.language in ["javascript", "typescript"]:
+                return f"process.env['{key}'] = '{value}'"
+            elif self.language == "deno":
+                return f"Deno.env.set('{key}', '{value}')"
+            elif self.language == "r":
+                return f'Sys.setenv({key} = "{value}")'
+            elif self.language == "java":
+                return f'System.setProperty("{key}", "{value}");'
+            elif self.language == "bash":
+                return f"export {key}='{value}'"
+        elif action == "delete":
+            if self.language == "python":
+                return f"import os; del os.environ['{key}']"
+            elif self.language in ["javascript", "typescript"]:
+                return f"delete process.env['{key}']"
+            elif self.language == "deno":
+                return f"Deno.env.delete('{key}')"
+            elif self.language == "r":
+                return f"Sys.unsetenv('{key}')"
+            elif self.language == "java":
+                return f'System.clearProperty("{key}");'
+            elif self.language == "bash":
+                return f"unset {key}"
+        return ""
+
+    def _build_env_vars_code(self, env_vars: Dict[StrictStr, str]) -> str:
+        """Build environment variable code for the current language."""
+        env_commands = []
+        for k, v in env_vars.items():
+            command = self._get_env_var_command(k, v, "set")
+            if command:
+                env_commands.append(command)
+        
+        return "\n".join(env_commands)
+
+    def _build_env_vars_cleanup_code(self, env_vars: Dict[StrictStr, str]) -> str:
+        """Build environment variable cleanup code for the current language."""
+        cleanup_commands = []
+        
+        for key in env_vars:
+            # Check if this var exists in global env vars
+            if self.global_env_vars and key in self.global_env_vars:
+                # Reset to global value
+                value = self.global_env_vars[key]
+                command = self._get_env_var_command(key, value, "set")
+            else:
+                # Remove the variable
+                command = self._get_env_var_command(key, "", "delete")
+            
+            if command:
+                cleanup_commands.append(command)
+        
+        return "\n".join(cleanup_commands)
+
     async def _wait_for_result(self, message_id: str):
         queue = self._executions[message_id].queue
 
@@ -133,84 +192,6 @@ 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]):
-        message_id = str(uuid.uuid4())
-        self._executions[message_id] = Execution(in_background=True)
-
-        env_commands = []
-        for k, v in env_vars.items():
-            if self.language == "python":
-                env_commands.append(f"import os; 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}'")
-            else:
-                return
-
-        if env_commands:
-            env_vars_snippet = "\n".join(env_commands)
-            logger.info(f"Setting env vars: {env_vars_snippet} for {self.language}")
-            request = self._get_execute_request(message_id, env_vars_snippet, True)
-            await self._ws.send(request)
-
-            async for item in self._wait_for_result(message_id):
-                if item["type"] == "error":
-                    raise ExecutionError(f"Error during execution: {item}")
-
-    async def reset_env_vars(self, env_vars: Dict[StrictStr, str]):
-        # 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 self.global_env_vars and key in self.global_env_vars:
-                vars_to_reset[key] = self.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:
-            message_id = str(uuid.uuid4())
-            self._executions[message_id] = Execution(in_background=True)
-
-            remove_commands = []
-            for key in vars_to_remove:
-                if self.language == "python":
-                    remove_commands.append(f"import os; 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}")
-                else:
-                    return
-            
-            if remove_commands:
-                remove_snippet = "\n".join(remove_commands)
-                logger.info(f"Removing env vars: {remove_snippet} for {self.language}")
-                request = self._get_execute_request(message_id, remove_snippet, True)
-                await self._ws.send(request)
-
-                async for item in self._wait_for_result(message_id):
-                    if item["type"] == "error":
-                        raise ExecutionError(f"Error during execution: {item}")
-
     async def change_current_directory(
         self, path: Union[str, StrictStr], language: str
     ):
@@ -256,12 +237,22 @@ async def execute(
             raise Exception("WebSocket not connected")
 
         async with self._lock:
-            # set env vars (will override global env vars)
+            # Build the complete code snippet with env vars
+            complete_code = code
+            
             if env_vars:
-                await self.set_env_vars(env_vars)
-
-            logger.info(code)
-            request = self._get_execute_request(message_id, code, False)
+                # Add env var setup at the beginning
+                env_setup_code = self._build_env_vars_code(env_vars)
+                if env_setup_code:
+                    complete_code = f"{env_setup_code}\n{complete_code}"
+                
+                # Add env var cleanup at the end
+                env_cleanup_code = self._build_env_vars_cleanup_code(env_vars)
+                if env_cleanup_code:
+                    complete_code = f"{complete_code}\n{env_cleanup_code}"
+
+            logger.info(f"Executing complete code: {complete_code}")
+            request = self._get_execute_request(message_id, complete_code, False)
 
             # Send the code for execution
             await self._ws.send(request)
@@ -272,10 +263,6 @@ 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")

From 0f3766a884957802d7831b2a135cf9c466fc0ab0 Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Wed, 20 Aug 2025 15:27:36 +0200
Subject: [PATCH 02/12] split env var snippet into two different methods

---
 template/server/messaging.py | 72 ++++++++++++++++++------------------
 1 file changed, 37 insertions(+), 35 deletions(-)

diff --git a/template/server/messaging.py b/template/server/messaging.py
index c6f6f222..e578d12b 100644
--- a/template/server/messaging.py
+++ b/template/server/messaging.py
@@ -114,47 +114,49 @@ def _get_execute_request(
             }
         )
 
-    def _get_env_var_command(self, key: str, value: str, action: str = "set") -> str:
-        """Get environment variable command for the current language."""
-        if action == "set":
-            if self.language == "python":
-                return f"import os; os.environ['{key}'] = '{value}'"
-            elif self.language in ["javascript", "typescript"]:
-                return f"process.env['{key}'] = '{value}'"
-            elif self.language == "deno":
-                return f"Deno.env.set('{key}', '{value}')"
-            elif self.language == "r":
-                return f'Sys.setenv({key} = "{value}")'
-            elif self.language == "java":
-                return f'System.setProperty("{key}", "{value}");'
-            elif self.language == "bash":
-                return f"export {key}='{value}'"
-        elif action == "delete":
-            if self.language == "python":
-                return f"import os; del os.environ['{key}']"
-            elif self.language in ["javascript", "typescript"]:
-                return f"delete process.env['{key}']"
-            elif self.language == "deno":
-                return f"Deno.env.delete('{key}')"
-            elif self.language == "r":
-                return f"Sys.unsetenv('{key}')"
-            elif self.language == "java":
-                return f'System.clearProperty("{key}");'
-            elif self.language == "bash":
-                return f"unset {key}"
+    def _set_env_var_snippet(self, key: str, value: str) -> str:
+        """Get environment variable set command for the current language."""
+        if self.language == "python":
+            return f"import os; os.environ['{key}'] = '{value}'"
+        elif self.language in ["javascript", "typescript"]:
+            return f"process.env['{key}'] = '{value}'"
+        elif self.language == "deno":
+            return f"Deno.env.set('{key}', '{value}')"
+        elif self.language == "r":
+            return f'Sys.setenv({key} = "{value}")'
+        elif self.language == "java":
+            return f'System.setProperty("{key}", "{value}");'
+        elif self.language == "bash":
+            return f"export {key}='{value}'"
         return ""
 
-    def _build_env_vars_code(self, env_vars: Dict[StrictStr, str]) -> str:
+    def _delete_env_var_snippet(self, key: str) -> str:
+        """Get environment variable delete command for the current language."""
+        if self.language == "python":
+            return f"import os; del os.environ['{key}']"
+        elif self.language in ["javascript", "typescript"]:
+            return f"delete process.env['{key}']"
+        elif self.language == "deno":
+            return f"Deno.env.delete('{key}')"
+        elif self.language == "r":
+            return f"Sys.unsetenv('{key}')"
+        elif self.language == "java":
+            return f'System.clearProperty("{key}");'
+        elif self.language == "bash":
+            return f"unset {key}"
+        return ""
+
+    def _set_env_vars_code(self, env_vars: Dict[StrictStr, str]) -> str:
         """Build environment variable code for the current language."""
         env_commands = []
         for k, v in env_vars.items():
-            command = self._get_env_var_command(k, v, "set")
+            command = self._set_env_var_snippet(k, v)
             if command:
                 env_commands.append(command)
         
         return "\n".join(env_commands)
 
-    def _build_env_vars_cleanup_code(self, env_vars: Dict[StrictStr, str]) -> str:
+    def _reset_env_vars_code(self, env_vars: Dict[StrictStr, str]) -> str:
         """Build environment variable cleanup code for the current language."""
         cleanup_commands = []
         
@@ -163,10 +165,10 @@ def _build_env_vars_cleanup_code(self, env_vars: Dict[StrictStr, str]) -> str:
             if self.global_env_vars and key in self.global_env_vars:
                 # Reset to global value
                 value = self.global_env_vars[key]
-                command = self._get_env_var_command(key, value, "set")
+                command = self._set_env_var_snippet(key, value)
             else:
                 # Remove the variable
-                command = self._get_env_var_command(key, "", "delete")
+                command = self._delete_env_var_snippet(key)
             
             if command:
                 cleanup_commands.append(command)
@@ -242,12 +244,12 @@ async def execute(
             
             if env_vars:
                 # Add env var setup at the beginning
-                env_setup_code = self._build_env_vars_code(env_vars)
+                env_setup_code = self._set_env_vars_code(env_vars)
                 if env_setup_code:
                     complete_code = f"{env_setup_code}\n{complete_code}"
                 
                 # Add env var cleanup at the end
-                env_cleanup_code = self._build_env_vars_cleanup_code(env_vars)
+                env_cleanup_code = self._reset_env_vars_code(env_vars)
                 if env_cleanup_code:
                     complete_code = f"{complete_code}\n{env_cleanup_code}"
 

From fbb037e8d70c2758c13b8f219e4f4961663daff4 Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Wed, 20 Aug 2025 15:55:16 +0200
Subject: [PATCH 03/12] handle setting global env vars (remove in background)

---
 template/server/main.py      |  5 -----
 template/server/messaging.py | 35 ++++++++++++++++++++++++++++++-----
 2 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/template/server/main.py b/template/server/main.py
index b353d514..7c59076d 100644
--- a/template/server/main.py
+++ b/template/server/main.py
@@ -105,11 +105,6 @@ async def post_execute(request: ExecutionRequest):
             status_code=404,
         )
 
-    # set global env vars if not set on first execution
-    if not ws.global_env_vars:
-        ws.global_env_vars = await get_envs()
-        await ws.set_env_vars(ws.global_env_vars)
-
     return StreamingListJsonResponse(
         ws.execute(
             request.code,
diff --git a/template/server/messaging.py b/template/server/messaging.py
index e578d12b..2e48551e 100644
--- a/template/server/messaging.py
+++ b/template/server/messaging.py
@@ -23,6 +23,7 @@
     UnexpectedEndOfExecution,
 )
 from errors import ExecutionError
+from envs import get_envs
 
 logger = logging.getLogger(__name__)
 
@@ -48,6 +49,7 @@ class ContextWebSocket:
     _ws: Optional[WebSocketClientProtocol] = None
     _receive_task: Optional[asyncio.Task] = None
     global_env_vars: Optional[Dict[StrictStr, str]] = None
+    global_env_vars_set = False
 
     def __init__(
         self,
@@ -175,6 +177,23 @@ def _reset_env_vars_code(self, env_vars: Dict[StrictStr, str]) -> str:
         
         return "\n".join(cleanup_commands)
 
+    async def _cleanup_env_vars(self, env_vars: Dict[StrictStr, str]):
+        """Clean up environment variables in a separate execution request."""
+        message_id = str(uuid.uuid4())
+        self._executions[message_id] = Execution(in_background=True)
+
+        cleanup_code = self._reset_env_vars_code(env_vars)
+        if cleanup_code:
+            logger.info(f"Cleaning up env vars: {cleanup_code}")
+            request = self._get_execute_request(message_id, cleanup_code, True)
+            await self._ws.send(request)
+
+            async for item in self._wait_for_result(message_id):
+                if item["type"] == "error":
+                    logger.error(f"Error during env var cleanup: {item}")
+
+        del self._executions[message_id]
+
     async def _wait_for_result(self, message_id: str):
         queue = self._executions[message_id].queue
 
@@ -241,17 +260,19 @@ async def execute(
         async with self._lock:
             # Build the complete code snippet with env vars
             complete_code = code
+
+            if not self.global_env_vars:
+                self.global_env_vars = await get_envs()
+
+            if not self.global_env_vars_set and self.global_env_vars:
+                complete_code = self._set_env_vars_code(self.global_env_vars)
+                self.global_env_vars_set = True
             
             if env_vars:
                 # Add env var setup at the beginning
                 env_setup_code = self._set_env_vars_code(env_vars)
                 if env_setup_code:
                     complete_code = f"{env_setup_code}\n{complete_code}"
-                
-                # Add env var cleanup at the end
-                env_cleanup_code = self._reset_env_vars_code(env_vars)
-                if env_cleanup_code:
-                    complete_code = f"{complete_code}\n{env_cleanup_code}"
 
             logger.info(f"Executing complete code: {complete_code}")
             request = self._get_execute_request(message_id, complete_code, False)
@@ -265,6 +286,10 @@ async def execute(
 
             del self._executions[message_id]
 
+        # Clean up env vars in a separate request after the main code has run (outside the lock)
+        if env_vars:
+            asyncio.create_task(self._cleanup_env_vars(env_vars))
+
     async def _receive_message(self):
         if not self._ws:
             logger.error("No WebSocket connection")

From 775725160ad07af33c93e98f709e24c0dfddcb3f Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Wed, 20 Aug 2025 16:10:52 +0200
Subject: [PATCH 04/12] cursor suggestion

---
 template/server/messaging.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/template/server/messaging.py b/template/server/messaging.py
index 2e48551e..337f322d 100644
--- a/template/server/messaging.py
+++ b/template/server/messaging.py
@@ -265,7 +265,7 @@ async def execute(
                 self.global_env_vars = await get_envs()
 
             if not self.global_env_vars_set and self.global_env_vars:
-                complete_code = self._set_env_vars_code(self.global_env_vars)
+                complete_code = f"{self._set_env_vars_code(self.global_env_vars)}\n{complete_code}"
                 self.global_env_vars_set = True
             
             if env_vars:

From aef5740c9ba269b9f8ed6b76e7530739af6c4b87 Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Wed, 20 Aug 2025 17:59:10 +0200
Subject: [PATCH 05/12] fix identation issues for env blocks

---
 python/tests/conftest.py     |  4 ++--
 template/server/messaging.py | 41 +++++++++++++++++++++++++++++++++---
 2 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/python/tests/conftest.py b/python/tests/conftest.py
index 30276cb1..daac9822 100644
--- a/python/tests/conftest.py
+++ b/python/tests/conftest.py
@@ -17,7 +17,7 @@ def template():
 
 @pytest.fixture()
 def sandbox(template, debug):
-    sandbox = Sandbox(template, timeout=timeout)
+    sandbox = Sandbox(template, timeout=timeout, debug=debug)
 
     try:
         yield sandbox
@@ -33,7 +33,7 @@ def sandbox(template, debug):
 
 @pytest_asyncio.fixture
 async def async_sandbox(template, debug):
-    async_sandbox = await AsyncSandbox.create(template, timeout=timeout)
+    async_sandbox = await AsyncSandbox.create(template, timeout=timeout, debug=debug)
 
     try:
         yield async_sandbox
diff --git a/template/server/messaging.py b/template/server/messaging.py
index 337f322d..a4207438 100644
--- a/template/server/messaging.py
+++ b/template/server/messaging.py
@@ -3,6 +3,7 @@
 import logging
 import uuid
 import asyncio
+import textwrap
 
 from asyncio import Queue
 from typing import (
@@ -177,6 +178,32 @@ def _reset_env_vars_code(self, env_vars: Dict[StrictStr, str]) -> str:
         
         return "\n".join(cleanup_commands)
 
+    def _get_code_indentation(self, code: str) -> str:
+        """Get the indentation from the first non-empty line of code."""
+        if not code or not code.strip():
+            return ""
+        
+        lines = code.split('\n')
+        for line in lines:
+            if line.strip():  # First non-empty line
+                return line[:len(line) - len(line.lstrip())]
+        
+        return ""
+
+    def _indent_code_with_level(self, code: str, indent_level: str) -> str:
+        """Apply the given indentation level to each line of code."""
+        if not code or not indent_level:
+            return code
+        
+        lines = code.split('\n')
+        indented_lines = []
+        
+        for line in lines:
+            if line.strip():  # Non-empty lines
+                indented_lines.append(indent_level + line)
+        
+        return '\n'.join(indented_lines)
+
     async def _cleanup_env_vars(self, env_vars: Dict[StrictStr, str]):
         """Clean up environment variables in a separate execution request."""
         message_id = str(uuid.uuid4())
@@ -258,6 +285,9 @@ async def execute(
             raise Exception("WebSocket not connected")
 
         async with self._lock:
+            # Get the indentation level from the code
+            code_indent = self._get_code_indentation(code)
+            
             # Build the complete code snippet with env vars
             complete_code = code
 
@@ -265,14 +295,18 @@ async def execute(
                 self.global_env_vars = await get_envs()
 
             if not self.global_env_vars_set and self.global_env_vars:
-                complete_code = f"{self._set_env_vars_code(self.global_env_vars)}\n{complete_code}"
+                env_setup_code = self._set_env_vars_code(self.global_env_vars)
+                if env_setup_code:
+                    indented_env_code = self._indent_code_with_level(env_setup_code, code_indent)
+                    complete_code = f"{indented_env_code}\n{complete_code}"
                 self.global_env_vars_set = True
             
             if env_vars:
                 # Add env var setup at the beginning
                 env_setup_code = self._set_env_vars_code(env_vars)
                 if env_setup_code:
-                    complete_code = f"{env_setup_code}\n{complete_code}"
+                    indented_env_code = self._indent_code_with_level(env_setup_code, code_indent)
+                    complete_code = f"{indented_env_code}\n{complete_code}"
 
             logger.info(f"Executing complete code: {complete_code}")
             request = self._get_execute_request(message_id, complete_code, False)
@@ -448,7 +482,8 @@ async def close(self):
         if self._ws is not None:
             await self._ws.close()
 
-        self._receive_task.cancel()
+        if self._receive_task is not None:
+            self._receive_task.cancel()
 
         for execution in self._executions.values():
             execution.queue.put_nowait(UnexpectedEndOfExecution())

From 37c2a2286a54dd37913d9af9b77cbb17d5286718 Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Wed, 20 Aug 2025 18:39:17 +0200
Subject: [PATCH 06/12] fix a race condition when subsequent request fire too
 rapidly before the cleanup could complete

---
 template/server/messaging.py | 52 +++++++++++++++++++++++++-----------
 1 file changed, 37 insertions(+), 15 deletions(-)

diff --git a/template/server/messaging.py b/template/server/messaging.py
index a4207438..19946c9c 100644
--- a/template/server/messaging.py
+++ b/template/server/messaging.py
@@ -3,7 +3,6 @@
 import logging
 import uuid
 import asyncio
-import textwrap
 
 from asyncio import Queue
 from typing import (
@@ -51,6 +50,7 @@ class ContextWebSocket:
     _receive_task: Optional[asyncio.Task] = None
     global_env_vars: Optional[Dict[StrictStr, str]] = None
     global_env_vars_set = False
+    _cleanup_task: Optional[asyncio.Task] = None
 
     def __init__(
         self,
@@ -209,17 +209,21 @@ async def _cleanup_env_vars(self, env_vars: Dict[StrictStr, str]):
         message_id = str(uuid.uuid4())
         self._executions[message_id] = Execution(in_background=True)
 
-        cleanup_code = self._reset_env_vars_code(env_vars)
-        if cleanup_code:
-            logger.info(f"Cleaning up env vars: {cleanup_code}")
-            request = self._get_execute_request(message_id, cleanup_code, True)
-            await self._ws.send(request)
-
-            async for item in self._wait_for_result(message_id):
-                if item["type"] == "error":
-                    logger.error(f"Error during env var cleanup: {item}")
-
-        del self._executions[message_id]
+        try:
+            cleanup_code = self._reset_env_vars_code(env_vars)
+            if cleanup_code:
+                logger.info(f"Cleaning up env vars: {cleanup_code}")
+                request = self._get_execute_request(message_id, cleanup_code, True)
+                await self._ws.send(request)
+
+                async for item in self._wait_for_result(message_id):
+                    if item["type"] == "error":
+                        logger.error(f"Error during env var cleanup: {item}")
+        finally:
+            del self._executions[message_id]
+            # Clear the task reference when cleanup is complete
+            if self._cleanup_task and self._cleanup_task.done():
+                self._cleanup_task = None
 
     async def _wait_for_result(self, message_id: str):
         queue = self._executions[message_id].queue
@@ -285,6 +289,16 @@ async def execute(
             raise Exception("WebSocket not connected")
 
         async with self._lock:
+            # Wait for any pending cleanup task to complete
+            if self._cleanup_task and not self._cleanup_task.done():
+                logger.debug("Waiting for pending cleanup task to complete")
+                try:
+                    await self._cleanup_task
+                except Exception as e:
+                    logger.warning(f"Cleanup task failed: {e}")
+                finally:
+                    self._cleanup_task = None
+            
             # Get the indentation level from the code
             code_indent = self._get_code_indentation(code)
             
@@ -320,9 +334,9 @@ async def execute(
 
             del self._executions[message_id]
 
-        # Clean up env vars in a separate request after the main code has run (outside the lock)
-        if env_vars:
-            asyncio.create_task(self._cleanup_env_vars(env_vars))
+            # Clean up env vars in a separate request after the main code has run
+            if env_vars:
+                self._cleanup_task = asyncio.create_task(self._cleanup_env_vars(env_vars))
 
     async def _receive_message(self):
         if not self._ws:
@@ -485,5 +499,13 @@ async def close(self):
         if self._receive_task is not None:
             self._receive_task.cancel()
 
+        # Cancel any pending cleanup task
+        if self._cleanup_task and not self._cleanup_task.done():
+            self._cleanup_task.cancel()
+            try:
+                await self._cleanup_task
+            except asyncio.CancelledError:
+                pass
+
         for execution in self._executions.values():
             execution.queue.put_nowait(UnexpectedEndOfExecution())

From b0f4da8d79367f33952ff1adbf4395fd1210ffbb Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Wed, 20 Aug 2025 18:47:13 +0200
Subject: [PATCH 07/12] make some ContextWebSocket private

---
 template/server/messaging.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/template/server/messaging.py b/template/server/messaging.py
index 19946c9c..f14c3124 100644
--- a/template/server/messaging.py
+++ b/template/server/messaging.py
@@ -48,8 +48,8 @@ def __init__(self, in_background: bool = False):
 class ContextWebSocket:
     _ws: Optional[WebSocketClientProtocol] = None
     _receive_task: Optional[asyncio.Task] = None
-    global_env_vars: Optional[Dict[StrictStr, str]] = None
-    global_env_vars_set = False
+    _global_env_vars: Optional[Dict[StrictStr, str]] = None
+    _global_env_vars_set = False
     _cleanup_task: Optional[asyncio.Task] = None
 
     def __init__(
@@ -165,9 +165,9 @@ def _reset_env_vars_code(self, env_vars: Dict[StrictStr, str]) -> str:
         
         for key in env_vars:
             # Check if this var exists in global env vars
-            if self.global_env_vars and key in self.global_env_vars:
+            if self._global_env_vars and key in self._global_env_vars:
                 # Reset to global value
-                value = self.global_env_vars[key]
+                value = self._global_env_vars[key]
                 command = self._set_env_var_snippet(key, value)
             else:
                 # Remove the variable
@@ -305,11 +305,11 @@ async def execute(
             # Build the complete code snippet with env vars
             complete_code = code
 
-            if not self.global_env_vars:
-                self.global_env_vars = await get_envs()
+            if not self._global_env_vars:
+                self._global_env_vars = await get_envs()
 
-            if not self.global_env_vars_set and self.global_env_vars:
-                env_setup_code = self._set_env_vars_code(self.global_env_vars)
+            if not self._global_env_vars_set and self._global_env_vars:
+                env_setup_code = self._set_env_vars_code(self._global_env_vars)
                 if env_setup_code:
                     indented_env_code = self._indent_code_with_level(env_setup_code, code_indent)
                     complete_code = f"{indented_env_code}\n{complete_code}"

From 475d787128d36d25a118b9c8ea57c97d19df4c22 Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Wed, 20 Aug 2025 19:03:29 +0200
Subject: [PATCH 08/12] fixes underscore var

---
 template/server/messaging.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/template/server/messaging.py b/template/server/messaging.py
index f14c3124..c61c4a2d 100644
--- a/template/server/messaging.py
+++ b/template/server/messaging.py
@@ -281,8 +281,6 @@ async def execute(
         env_vars: Dict[StrictStr, str] = None,
     ):
         message_id = str(uuid.uuid4())
-        logger.debug(f"Sending code for the execution ({message_id}): {code}")
-
         self._executions[message_id] = Execution()
 
         if self._ws is None:
@@ -313,7 +311,7 @@ async def execute(
                 if env_setup_code:
                     indented_env_code = self._indent_code_with_level(env_setup_code, code_indent)
                     complete_code = f"{indented_env_code}\n{complete_code}"
-                self.global_env_vars_set = True
+                self._global_env_vars_set = True
             
             if env_vars:
                 # Add env var setup at the beginning
@@ -322,7 +320,7 @@ async def execute(
                     indented_env_code = self._indent_code_with_level(env_setup_code, code_indent)
                     complete_code = f"{indented_env_code}\n{complete_code}"
 
-            logger.info(f"Executing complete code: {complete_code}")
+            logger.info(f"Sending code for the execution ({message_id}): {complete_code}")
             request = self._get_execute_request(message_id, complete_code, False)
 
             # Send the code for execution

From 1118dd605457910c0c206e6ba663f2670b399551 Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Wed, 20 Aug 2025 19:20:34 +0200
Subject: [PATCH 09/12] sequential env var snippets

---
 template/server/messaging.py | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/template/server/messaging.py b/template/server/messaging.py
index c61c4a2d..9944a821 100644
--- a/template/server/messaging.py
+++ b/template/server/messaging.py
@@ -302,23 +302,21 @@ async def execute(
             
             # Build the complete code snippet with env vars
             complete_code = code
+            env_var_snippets = []
 
             if not self._global_env_vars:
                 self._global_env_vars = await get_envs()
 
             if not self._global_env_vars_set and self._global_env_vars:
-                env_setup_code = self._set_env_vars_code(self._global_env_vars)
-                if env_setup_code:
-                    indented_env_code = self._indent_code_with_level(env_setup_code, code_indent)
-                    complete_code = f"{indented_env_code}\n{complete_code}"
+                env_var_snippets.append(self._set_env_vars_code(self._global_env_vars))
                 self._global_env_vars_set = True
             
             if env_vars:
-                # Add env var setup at the beginning
-                env_setup_code = self._set_env_vars_code(env_vars)
-                if env_setup_code:
-                    indented_env_code = self._indent_code_with_level(env_setup_code, code_indent)
-                    complete_code = f"{indented_env_code}\n{complete_code}"
+                env_var_snippets.append(self._set_env_vars_code(env_vars))
+
+            if env_var_snippets:
+                indented_env_code = self._indent_code_with_level("\n".join(env_var_snippets), code_indent)
+                complete_code = f"{indented_env_code}\n{complete_code}"
 
             logger.info(f"Sending code for the execution ({message_id}): {complete_code}")
             request = self._get_execute_request(message_id, complete_code, False)

From a4d5f6154cfdb2119317c8981be6dd2bb4c3a25e Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Wed, 20 Aug 2025 19:47:58 +0200
Subject: [PATCH 10/12] removed dead code (Cursor)

---
 template/server/messaging.py | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/template/server/messaging.py b/template/server/messaging.py
index 9944a821..deae4f87 100644
--- a/template/server/messaging.py
+++ b/template/server/messaging.py
@@ -221,9 +221,6 @@ async def _cleanup_env_vars(self, env_vars: Dict[StrictStr, str]):
                         logger.error(f"Error during env var cleanup: {item}")
         finally:
             del self._executions[message_id]
-            # Clear the task reference when cleanup is complete
-            if self._cleanup_task and self._cleanup_task.done():
-                self._cleanup_task = None
 
     async def _wait_for_result(self, message_id: str):
         queue = self._executions[message_id].queue

From 24756560633eb755dca0975153f98b6b185f4bc3 Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Wed, 20 Aug 2025 20:00:49 +0200
Subject: [PATCH 11/12] none check for global env vars

---
 template/server/messaging.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/template/server/messaging.py b/template/server/messaging.py
index deae4f87..ca220bf7 100644
--- a/template/server/messaging.py
+++ b/template/server/messaging.py
@@ -201,6 +201,8 @@ def _indent_code_with_level(self, code: str, indent_level: str) -> str:
         for line in lines:
             if line.strip():  # Non-empty lines
                 indented_lines.append(indent_level + line)
+            else:
+                indented_lines.append(line)
         
         return '\n'.join(indented_lines)
 
@@ -301,10 +303,10 @@ async def execute(
             complete_code = code
             env_var_snippets = []
 
-            if not self._global_env_vars:
+            if self._global_env_vars is None:
                 self._global_env_vars = await get_envs()
 
-            if not self._global_env_vars_set and self._global_env_vars:
+            if not self._global_env_vars_set:
                 env_var_snippets.append(self._set_env_vars_code(self._global_env_vars))
                 self._global_env_vars_set = True
             

From 6aa447e25d674ad61032e0fdbe944173426770fa Mon Sep 17 00:00:00 2001
From: Mish <10400064+mishushakov@users.noreply.github.com>
Date: Thu, 21 Aug 2025 19:25:16 +0200
Subject: [PATCH 12/12] changed env var snippet logic and self._global_env_vars
 check

---
 template/server/messaging.py | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/template/server/messaging.py b/template/server/messaging.py
index ca220bf7..7c295b2d 100644
--- a/template/server/messaging.py
+++ b/template/server/messaging.py
@@ -49,7 +49,6 @@ class ContextWebSocket:
     _ws: Optional[WebSocketClientProtocol] = None
     _receive_task: Optional[asyncio.Task] = None
     _global_env_vars: Optional[Dict[StrictStr, str]] = None
-    _global_env_vars_set = False
     _cleanup_task: Optional[asyncio.Task] = None
 
     def __init__(
@@ -301,20 +300,19 @@ async def execute(
             
             # Build the complete code snippet with env vars
             complete_code = code
-            env_var_snippets = []
+            
+            global_env_vars_snippet = ""
+            env_vars_snippet = ""
 
             if self._global_env_vars is None:
                 self._global_env_vars = await get_envs()
-
-            if not self._global_env_vars_set:
-                env_var_snippets.append(self._set_env_vars_code(self._global_env_vars))
-                self._global_env_vars_set = True
+                global_env_vars_snippet = self._set_env_vars_code(self._global_env_vars)
             
             if env_vars:
-                env_var_snippets.append(self._set_env_vars_code(env_vars))
+                env_vars_snippet = self._set_env_vars_code(env_vars)
 
-            if env_var_snippets:
-                indented_env_code = self._indent_code_with_level("\n".join(env_var_snippets), code_indent)
+            if global_env_vars_snippet or env_vars_snippet:
+                indented_env_code = self._indent_code_with_level(f"{global_env_vars_snippet}\n{env_vars_snippet}", code_indent)
                 complete_code = f"{indented_env_code}\n{complete_code}"
 
             logger.info(f"Sending code for the execution ({message_id}): {complete_code}")