Merge pull request #35 from microsoft/kkaitepalli/exitcode-correction

Limit ShellCommunicator to bash

Author: KavyaSree2610

src/microbots/environment/local_docker/LocalDockerEnvironment.py

@@ -180,7 +180,27 @@ def copy_to_container(self, src_path: str, dest_path: str) -> bool:
             if not os.path.exists(src_path):
                 logger.error("❌ Source path does not exist: %s", src_path)
                 return False
-                
+            # Ensure destination directory exists inside container
+            dest_dir = os.path.dirname(dest_path)
+            if dest_dir and dest_dir != '/':
+                # Check if directory exists inside the container first
+                check_cmd = f"test -d {shlex.quote(dest_dir)}"
+                check_result = self.execute(check_cmd)
+
+                if check_result.return_code != 0:
+                    logger.debug("📁 Creating destination directory inside container: %s", dest_dir)
+                    mkdir_cmd = f"mkdir -p {shlex.quote(dest_dir)}"
+                    mkdir_result = self.execute(mkdir_cmd)
+ 
+                    if mkdir_result.return_code != 0:
+                        logger.error("❌ Failed to create destination directory %s: %s", 
+                                   dest_dir, mkdir_result.stderr)
+                        return False
+                    else:
+                        logger.debug("✅ Destination directory created: %s", dest_dir)
+                else:
+                    logger.debug("✅ Destination directory already exists: %s", dest_dir)
+
             # Use docker cp command to copy files/folders
             # Escape paths for shell safety
             src_escaped = shlex.quote(src_path)
@@ -189,7 +209,7 @@ def copy_to_container(self, src_path: str, dest_path: str) -> bool:
             # Build docker cp command
             cmd = f"docker cp {src_escaped} {self.container.id}:{dest_escaped}"
 
-            logger.info("📁 Copying %s to container:%s", src_path, dest_path)
+            logger.debug("📁 Copying %s to container:%s", src_path, dest_path)
 
             # Execute the copy command
             result = subprocess.run(
@@ -230,14 +250,28 @@ def copy_from_container(self, src_path: str, dest_path: str) -> bool:
             return False
 
         try:
+            # Check if source path exists inside the container
+            check_cmd = f"test -e {shlex.quote(src_path)}"
+            check_result = self.execute(check_cmd)
+
+            if check_result.return_code != 0:
+                logger.error("❌ Source path does not exist in container: %s", src_path)
+                return False
+
+            # Check if destination directory exists on host machine
+            dest_dir = os.path.dirname(dest_path)
+            if not os.path.exists(dest_dir):
+                logger.error("❌ Destination directory does not exist on host: %s", dest_dir)
+                return False
+
             # Escape paths for shell safety
             src_escaped = shlex.quote(src_path)
             dest_escaped = shlex.quote(dest_path)
 
             # Build docker cp command
             cmd = f"docker cp {self.container.id}:{src_escaped} {dest_escaped}"
 
-            logger.info("📁 Copying container:%s to %s", src_path, dest_path)
+            logger.debug("📁 Copying container:%s to %s", src_path, dest_path)
 
             # Execute the copy command
             result = subprocess.run(
@@ -249,7 +283,7 @@ def copy_from_container(self, src_path: str, dest_path: str) -> bool:
             )
 
             if result.returncode == 0:
-                logger.info("✅ Successfully copied container:%s to %s", src_path, dest_path)
+                logger.info("✅ Successfully copied from container:%s to %s", src_path, dest_path)
                 return True
             else:
                 logger.error("❌ Failed to copy file. Error: %s", result.stderr)

src/microbots/environment/local_docker/image_builder/ShellCommunicator.py

@@ -34,10 +34,12 @@ def __init__(self, shell_type: str = "bash", encoding: str = "utf-8"):
         Initialize the shell communicator.
 
         Args:
-            shell_type: Type of shell ("powershell", "cmd", "bash", "python")
+            shell_type: Type of shell (only "bash" is supported)
             encoding: Text encoding for communication
         """
         self.shell_type = shell_type.lower()
+        if self.shell_type != "bash":
+            raise ValueError(f"Unsupported shell type: {shell_type}. Only 'bash' is supported.")
         self.encoding = encoding
         self.process: Optional[subprocess.Popen] = None
         self.output_queue = queue.Queue()
@@ -47,13 +49,9 @@ def __init__(self, shell_type: str = "bash", encoding: str = "utf-8"):
         self.error_thread: Optional[threading.Thread] = None
         self.output_callback: Optional[Callable] = None
 
-        # Define shell commands
+        # Define shell commands - only supporting bash
         self.shell_commands = {
-            "powershell": ["powershell.exe", "-NoLogo", "-NoExit"],
-            "cmd": ["cmd.exe", "/k"],
             "bash": ["bash"],
-            "python": [sys.executable, "-i"],
-            "wsl": ["wsl.exe"],
         }
 
     def start_session(self) -> bool:
@@ -134,7 +132,6 @@ def _re_escape(self, command: str) -> str:
         # command = command.replace("&lt;", "<").replace("&gt;", ">")
         return command
 
-   # TODO: Exit code is not properly captured. Need to fix it.
     def send_command(
         self, command: str, wait_for_output: bool = True, timeout: float = 300
     ) -> CmdReturn:
@@ -155,28 +152,24 @@ def send_command(
 
         try:
             command = self._re_escape(command)
-            # Send the command
-            self.process.stdin.write(command + "\n")
-            self.process.stdin.flush()
-            logger.debug("➡️ Sent command: %s", command)
-
+            
             if not wait_for_output:
+                # Send the command without marker for async execution
+                self.process.stdin.write(command + "\n")
+                self.process.stdin.flush()
+                logger.debug("➡️ Sent async command: %s", command)
                 return CmdReturn(stdout="ASYNC: Not waiting for completion", stderr="", return_code=0)
 
             # Generate a unique command completion marker
             marker = f"__COMMAND_COMPLETE_{int(time.time() * 1000000)}__"
 
-            # Send the marker command based on shell type
-            if self.shell_type in ["bash", "wsl"]:
-                self.process.stdin.write(f"echo '{marker}'; echo $? > /tmp/last_exit_code\n")
-            elif self.shell_type == "powershell":
-                self.process.stdin.write(f"echo '{marker}'; echo $LASTEXITCODE\n")
-            elif self.shell_type == "cmd":
-                self.process.stdin.write(f"echo {marker} & echo %ERRORLEVEL%\n")
-            elif self.shell_type == "python":
-                self.process.stdin.write(f"print('{marker}')\n")
-
+            # For bash only: Send command + marker in a single line to capture correct exit code
+            combined_command = f"{command}; echo '{marker}' $?"
+            
+            # Send the combined command
+            self.process.stdin.write(combined_command + "\n")
             self.process.stdin.flush()
+            logger.debug("➡️ Sent command: %s", command)
 
             # Collect output until marker is found or timeout
             output_lines = []
@@ -193,23 +186,24 @@ def send_command(
                     # Check if this is our completion marker
                     if marker in line:
                         marker_found = True
-                        # For bash/wsl, try to get the exit code from the next line
-                        if self.shell_type in ["bash", "wsl"]:
-                            try:
-                                # Try to get exit code from next output
-                                stream_type, exit_code_line = self.output_queue.get(timeout=0.5)
-                                if exit_code_line.strip().isdigit():
-                                    last_exit_code = int(exit_code_line.strip())
-                            except (queue.Empty, ValueError):
-                                pass
-                        elif self.shell_type in ["powershell", "cmd"]:
-                            try:
-                                # Try to get exit code from next output
-                                stream_type, exit_code_line = self.output_queue.get(timeout=0.5)
-                                if exit_code_line.strip().isdigit():
-                                    last_exit_code = int(exit_code_line.strip())
-                            except (queue.Empty, ValueError):
-                                pass
+                        # For bash, the exit code is on the same line after the marker
+                        try:
+                            # Extract exit code from the same line as the marker
+                            # Format: "__COMMAND_COMPLETE_xxxxx__ exit_code"
+                            parts = line.split()
+                            if len(parts) >= 2:
+                                exit_code_str = parts[-1].strip()
+                                # Handle bash exit codes (0-255 only)
+                                if exit_code_str.isdigit():
+                                    last_exit_code = int(exit_code_str)
+                                else:
+                                    last_exit_code = 1
+                            else:
+                                last_exit_code = 1
+                        except (ValueError, AttributeError, IndexError):
+                            # Default to 1 if parsing fails
+                            last_exit_code = 1
+                        logger.debug("🔍 Found marker with exit code: %s", last_exit_code)
                         continue
 
                     # Add output to appropriate list
@@ -236,7 +230,6 @@ def send_command(
                 except queue.Empty:
                     break
 
-            # TODO: Final return code is not correct. Need a fix
             final_return_code = last_exit_code if marker_found else (1 if error_lines else 0)
 
             # Handle timeout case
@@ -295,13 +288,8 @@ def close_session(self):
 
         if self.process:
             try:
-                # Try to terminate gracefully
-                if self.shell_type == "powershell":
-                    self.send_command("exit", wait_for_output=False)
-                elif self.shell_type == "cmd":
-                    self.send_command("exit", wait_for_output=False)
-                else:
-                    self.send_command("exit", wait_for_output=False)
+                # Try to terminate gracefully with bash exit command
+                self.send_command("exit", wait_for_output=False)
 
                 # Wait a bit for graceful shutdown
                 time.sleep(1)

src/microbots/environment/local_docker/image_builder/dockerShell.py

@@ -1,10 +1,24 @@
 import os
+import logging
 
 import uvicorn
 from fastapi import FastAPI
 from pydantic import BaseModel
 from ShellCommunicator import ShellCommunicator
 
+# Configure logging to see all logs including ShellCommunicator
+logging.basicConfig(
+    level=logging.DEBUG,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.StreamHandler()  # Ensure logs go to stdout
+    ]
+)
+
+# Set specific logger levels
+logging.getLogger('ShellCommunicator').setLevel(logging.DEBUG)
+logging.getLogger('uvicorn').setLevel(logging.INFO)
+
 shell = ShellCommunicator("bash")
 shell.start_session()
 

test/environment/local_docker/TestFileCopy.py

@@ -4,10 +4,11 @@
 """
 
 import os
-import tempfile
-import unittest
 import sys
-
+from pathlib import Path
+import logging
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
 # Add src directory to path to import from local source
 sys.path.insert(
     0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../src"))
@@ -25,8 +26,9 @@ def test_copy_file(self):
 
         try:
             # Copy to container
-            # Give absolute path
-            result = env.copy_to_container("/home/kkaitepalli/minions/README.md", "/var/log/README.md")
+            # Get path to countries.txt file specifically
+            countries_file_path = Path(__file__).parent.parent.parent / "bot" / "countries_to_capital" / "countries_dir" / "countries.txt"
+            result = env.copy_to_container(str(countries_file_path), "/var/log/")
 
             # Verify
             print(f"Copy result: {result}")
@@ -36,7 +38,8 @@ def test_copy_file(self):
                 print("❌ Copy failed")
 
             # Test copying from container to host
-            result_back = env.copy_from_container("/var/log/README.md", "/home/kkaitepalli/temp/README_copied.md")
+            # Use /tmp/ which is available and writable on all systems
+            result_back = env.copy_from_container("/var/log/countries.txt", "/tmp/")
             print(f"Copy back result: {result_back}")
             if result_back:
                 print("✅ Copy back succeeded")