Fix packaging, set TERM type correctly

code_execution_tool.py → code_execution_mcp/code_execution_tool.py

@@ -5,11 +5,10 @@
 import re
 import os
 import sys
-from helpers.print_style import PrintStyle
-from helpers.shell_local import LocalInteractiveSession
-from helpers.strings import truncate_text as truncate_text_string
-from helpers.log import Log
-
+from .helpers.print_style import PrintStyle
+from .helpers.shell_local import LocalInteractiveSession
+from .helpers.strings import truncate_text as truncate_text_string
+from .helpers.log import Log
 
 def truncate_text_agent(output: str, threshold: int = 1000000) -> str:
     if len(output) <= threshold:

helpers/shell_local.py → code_execution_mcp/helpers/shell_local.py

@@ -7,14 +7,25 @@
 from . import tty_session
 from .shell_utils import clean_string
 
+
 class LocalInteractiveSession:
-    def __init__(self, executable: str | None = None):
+    def __init__(
+        self,
+        executable: str | None = None,
+        term_type: str = "dumb",
+    ):
         self.executable = executable
-        self.session: tty_session.TTYSession|None = None
+        self.term_type = term_type
+        self.session: tty_session.TTYSession | None = None
         self.full_output = ''
 
     async def connect(self):
-        self.session = tty_session.TTYSession(self.executable, env=os.environ.copy())
+        env = os.environ.copy()
+        env["TERM"] = self.term_type
+        self.session = tty_session.TTYSession(
+            self.executable,
+            env=env,
+        )
         await self.session.start()
         await self.session.read_full_until_idle(idle_timeout=1, total_timeout=1)
 
@@ -28,7 +39,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")
@@ -46,4 +57,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

main.py → code_execution_mcp/main.py

@@ -2,7 +2,18 @@
 import os
 import sys
 from fastmcp import FastMCP
-from code_execution_tool import CodeExecutionTool
+
+try:
+    # Prefer package-style import when running as an installed package or with -m
+    from code_execution_mcp.code_execution_tool import CodeExecutionTool
+except ImportError:
+    # Fallback for running this file directly, e.g. `python /path/to/code_execution_mcp/main.py`
+    # Ensure the project root (parent of this directory) is on sys.path so that
+    # `code_execution_mcp` can be imported as a proper package and relative imports work.
+    package_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    if package_root not in sys.path:
+        sys.path.insert(0, package_root)
+    from code_execution_mcp.code_execution_tool import CodeExecutionTool
 
 # Helpers for debugging
 version = f"v0.1.3, Python {sys.version.split(' ')[0]}, executable={sys.executable}"

pyproject.toml

@@ -13,10 +13,14 @@ dependencies = [
 ]
 
 [project.scripts]
-code-execution-mcp = "main:main"
+code-execution-mcp = "code_execution_mcp.main:main"
 
 [tool.setuptools]
-py-modules = ["main", "code_execution_tool"]
+py-modules = ["code_execution_mcp.main", "code_execution_mcp.code_execution_tool"]
+include-package-data = true
+
+[tool.setuptools.package-data]
+"code_execution_mcp" = ["prompts/*.md"]
 
 [build-system]
 requires = ["setuptools>=61.0"]