feat: expose extra_env in LocalInteractiveSession and tool instance in tool_execute extensions

plugins/_code_execution/helpers/shell_local.py

@@ -1,3 +1,4 @@
+import os
 import platform
 import select
 import subprocess
@@ -8,14 +9,27 @@
 from plugins._code_execution.helpers import tty_session
 from plugins._code_execution.helpers.shell_ssh import clean_string
 
+# Environment variable keys that are safe to forward to the subprocess.
+# Deliberately excludes API keys, bearer tokens, secrets, and all framework
+# env vars — only the minimal set required for a functional interactive shell.
+_SAFE_ENV_KEYS = {"PATH", "HOME", "USER", "SHELL", "TERM", "LANG", "LC_ALL", "TMPDIR", "PWD"}
+
+
 class LocalInteractiveSession:
-    def __init__(self, cwd: str|None = None):
-        self.session: tty_session.TTYSession|None = None
+    def __init__(self, cwd: str | None = None, extra_env: dict | None = None):
+        self.session: tty_session.TTYSession | None = None
         self.full_output = ''
         self.cwd = cwd
+        self.extra_env = extra_env
 
     async def connect(self):
-        self.session = tty_session.TTYSession(runtime.get_terminal_executable(), cwd=self.cwd)
+        # When extra_env is provided, build a clean env from the safe whitelist
+        # only — never merge the full os.environ, which would expose API keys
+        # and other framework secrets to the subprocess.
+        env = (
+            {k: v for k, v in os.environ.items() if k in _SAFE_ENV_KEYS} | self.extra_env
+        ) if self.extra_env else None
+        self.session = tty_session.TTYSession(runtime.get_terminal_executable(), cwd=self.cwd, env=env)
         await self.session.start()
         await self.session.read_full_until_idle(idle_timeout=1, total_timeout=1)
 
@@ -29,7 +43,7 @@ async def send_command(self, command: str):
             raise Exception("Shell not connected")
         self.full_output = ""
         await self.session.sendline(command)
- 
+
     async def read_output(self, timeout: float = 0, reset_full_output: bool = False) -> Tuple[str, Optional[str]]:
         if not self.session:
             raise Exception("Shell not connected")
@@ -47,4 +61,4 @@ async def read_output(self, timeout: float = 0, reset_full_output: bool = False)
 
         if not partial_output:
             return clean_full_output, None
-        return clean_full_output, partial_output
+        return clean_full_output, partial_output

plugins/_code_execution/helpers/shell_ssh.py

@@ -1,5 +1,6 @@
 import asyncio
 import paramiko
+import shlex
 import time
 import re
 from typing import Tuple
@@ -14,7 +15,8 @@ class SSHInteractiveSession:
     # ps1_label = "SSHInteractiveSession CLI>"
 
     def __init__(
-        self, logger: Log, hostname: str, port: int, username: str, password: str, cwd: str|None = None
+        self, logger: Log, hostname: str, port: int, username: str, password: str,
+        cwd: str | None = None, extra_env: dict | None = None
     ):
         self.logger = logger
         self.hostname = hostname
@@ -28,6 +30,7 @@ def __init__(
         self.last_command = b""
         self.trimmed_command_length = 0  # Initialize trimmed_command_length
         self.cwd = cwd
+        self.extra_env = extra_env
 
     async def connect(self, keepalive_interval: int = 5):
         """
@@ -37,7 +40,7 @@ async def connect(self, keepalive_interval: int = 5):
         ----------
         keepalive_interval : int
             Interval in **seconds** between keep-alive packets sent by Paramiko.
-            A value ≤ 0 disables Paramiko’s keep-alive feature.
+            A value ≤ 0 disables Paramiko's keep-alive feature.
         """
         errors = 0
         while True:
@@ -66,6 +69,17 @@ async def connect(self, keepalive_interval: int = 5):
                 initial_command = "unset PROMPT_COMMAND PS0; stty -echo"
                 if self.cwd:
                     initial_command = f"cd {self.cwd}; {initial_command}"
+
+                # When extra_env is provided, prepend export statements so the
+                # variables are available for the entire session.  Values are
+                # shell-quoted via shlex.quote to prevent injection.
+                if self.extra_env:
+                    exports = "; ".join(
+                        f"export {k}={shlex.quote(str(v))}"
+                        for k, v in self.extra_env.items()
+                    )
+                    initial_command = f"{exports}; {initial_command}"
+
                 self.shell.send(f"{initial_command}\n".encode())
 
                 # wait for initial prompt/output to settle
@@ -104,7 +118,7 @@ async def send_command(self, command: str):
         self.last_command = command.encode()
         self.trimmed_command_length = 0
         self.shell.send(self.last_command)
-        
+
     async def read_output(
         self, timeout: float = 0, reset_full_output: bool = False
     ) -> Tuple[str, str]:
@@ -138,7 +152,7 @@ async def read_output(
             #         deviation_threshold=8,
             #         deviation_reset=2,
             #         ignore_patterns=[
-            #             rb"\x1b\[\?\d{4}[a-zA-Z](?:> )?",  # ANSI escape sequences
+            #             rb"\[\?\d{4}[a-zA-Z](?:> )?",  # ANSI escape sequences
             #             rb"\r",  # Carriage return
             #             rb">\s",  # Greater-than symbol
             #         ],
@@ -212,13 +226,14 @@ def recv_all(num_bytes):
 
         return data
 
+
 def clean_string(input_string):
     # Remove ANSI escape codes
-    ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
+    ansi_escape = re.compile(r"\x1b(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])")
     cleaned = ansi_escape.sub("", input_string)
 
     # remove null bytes
-    cleaned = cleaned.replace("\x00", "")
+    cleaned = cleaned.replace("", "")
 
     # remove ipython \r\r\n> sequences from the start
     cleaned = re.sub(r'^[ \r]*(?:\r*\n>[ \r]*)*', '', cleaned)