FROM feat/4-add-mcp-server TO development

.gitignore

@@ -2,4 +2,5 @@ uploads
 downloads
 toolkilt/.venv/
 toolkilt/__pycache__/
-toolkilt/tools/__pycache__/
+toolkilt/tools/__pycache__/
+**/__pycache__/

.vscode/launch.json

@@ -0,0 +1,19 @@
+{
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Python: MCP Server",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "${workspaceFolder}/api.py",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+            "env": {
+                "PYTHONPATH": "${workspaceFolder}",
+                // "FASTMCP_DEBUG": "true",
+                // "FASTMCP_LOG_LEVEL": "DEBUG",
+                "PYDEVD_DISABLE_FILE_VALIDATION": "1"
+            }
+        }
+    ]
+}

Dockerfile

@@ -1,11 +1,16 @@
 # Use the official Jupyter Notebook image
 FROM jupyter/base-notebook:latest
 
-# Install FastAPI and Uvicorn
-RUN pip install fastapi uvicorn python-multipart
-
 # Copy the API script to the root directory
+COPY mcp_wrap /home/jovyan/mcp_wrap
+COPY src /home/jovyan/src
 COPY api.py /home/jovyan/api.py
+COPY requirements.txt /home/jovyan/requirements.txt
+
+# Install FastAPI and Uvicorn
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir uv && \
+    uv pip install --system --no-cache-dir -r requirements.txt
 
 # Expose port 8888 for the Jupyter Notebook
 EXPOSE 8888

api.py

@@ -1,30 +1,42 @@
-from fastapi import FastAPI, HTTPException, File, UploadFile, Form
-from pydantic import BaseModel
-from fastapi.responses import FileResponse
+from fastapi import Body, Depends, FastAPI, HTTPException, File, Request, UploadFile, Form
+from fastapi.responses import FileResponse, JSONResponse
+from fastapi_mcp import FastApiMCP, AuthConfig
 import subprocess
 import os
-import shutil
-from typing import List, Dict
-
-app = FastAPI()
-
-class CodeExecutionRequest(BaseModel):
-    session_id: str
-    code: str
-    env: Dict[str, str] = {}
-
-class PackageInstallationRequest(BaseModel):
-    session_id: str
-    packages: List[str]
-
-class TerminateSessionRequest(BaseModel):
-    session_id: str
 
+from src.entities import *
+from src.controllers import ExecController
+from src.utils.session import get_session_id
+from src.config import Config
+from src.services.auth import authenticate_request
+TAGS = ['MCP']
+ID_NAME = "session_id"
+MCP_OPERATIONS = [
+    "install_packages",
+    "code_execute",
+    "terminate_session",
+]
+
+
+
+app = FastAPI(
+    title=Config.TITLE,
+    description=Config.DESCRIPTION,
+    version=Config.VERSION,
+    docs_url="/",
+    # redoc_url="/redoc"
+)
+exec_controller = ExecController()
 session_manager = {}
 
-@app.post("/install")
-def install_packages(request: PackageInstallationRequest):
-    session_id = request.session_id
+@app.post(
+    "/install", 
+    operation_id="install_packages",
+    description="Install packages in isolated environment",
+    tags=TAGS
+)
+def install_packages(request: PackageInstall):
+    session_id = get_session_id(request)
 
     if session_id not in session_manager:
         session_manager[session_id] = {
@@ -34,47 +46,73 @@ def install_packages(request: PackageInstallationRequest):
 
     try:
         # Install packages if any are provided and not already installed
-        for package in request.packages:
-            if package not in session_manager[session_id]["packages"]:
-                subprocess.check_call([f"pip install {package}"], shell=True)
-                session_manager[session_id]["packages"].add(package)
-
-        return {"status": "success", "installed_packages": list(session_manager[session_id]["packages"])}
+        return JSONResponse(
+            status_code=200, 
+            content={
+                ID_NAME: session_id,
+                **exec_controller.install_packages(session_id, session_manager, request.packages)
+            }
+        )
     except subprocess.CalledProcessError as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-@app.post("/execute")
-def run_code(request: CodeExecutionRequest):
-    session_id = request.session_id
-
+@app.post(
+    "/execute", 
+    operation_id="code_execute",
+    description="Execute code in isolated environment",
+    tags=TAGS,
+    dependencies=[Depends(authenticate_request)]
+)
+def code_execute(body: CodeExecution = Body(...)):
+    session_id = get_session_id(body)
     if session_id not in session_manager:
         session_manager[session_id] = {
             "packages": set(),
             "files": set()
         }
 
     try:
-        # Set environment variables
-        for key, value in request.env.items():
-            os.environ[key] = value
-
-        # Create session directory if it doesn't exist
-        session_dir = f"/tmp/{session_id}"
-        os.makedirs(session_dir, exist_ok=True)
-
         # Write code to a temporary file
-        code_file_path = f"{session_dir}/temp_code.py"
-        with open(code_file_path, "w") as code_file:
-            code_file.write(request.code)
+        code_file_path = exec_controller.exec_service.write_code(session_id, body.code)
         session_manager[session_id]["files"].add(code_file_path)
 
-        # Run the code
-        result = subprocess.run(["python", code_file_path], capture_output=True, text=True)
-
-        return {"status": "success", "output": result.stdout, "errors": result.stderr}
+        return JSONResponse(
+            status_code=200, 
+            content={
+                ID_NAME: session_id,
+                **exec_controller.execute_python(code_file_path).model_dump()
+            }
+        )
+
     except subprocess.CalledProcessError as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+@app.post(
+    "/terminate",
+    operation_id="terminate_session",
+    description="Terminate session",
+    tags=TAGS
+)
+def terminate_session(request: SessionId):
+    session_id = get_session_id(request)
+
+    if session_id not in session_manager:
+        raise HTTPException(status_code=404, detail="Session not found.")
+
+    try:
+        deleted_session = exec_controller.exec_service.uninstall_packages(session_id, session_manager)
+        return JSONResponse(
+            status_code=200, 
+            content={
+                ID_NAME: session_id,
+                **deleted_session
+            }
+        )
+    except subprocess.CalledProcessError as e:
+        raise HTTPException(status_code=500, detail=str(e))
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
 @app.post("/upload")
 async def create_upload_file(session_id: str = Form(...), file: UploadFile = File(...)):
     if session_id not in session_manager:
@@ -84,48 +122,14 @@ async def create_upload_file(session_id: str = Form(...), file: UploadFile = Fil
         }
 
     try:
-        # Create session directory if it doesn't exist
-        session_dir = f"/tmp/{session_id}"
-        os.makedirs(session_dir, exist_ok=True)
-
-        file_location = f"{session_dir}/{file.filename}"
-        with open(file_location, "wb+") as file_object:
-            shutil.copyfileobj(file.file, file_object)
-
-        session_manager[session_id]["files"].add(file_location)
-
-        return {"filename": file.filename, "location": file_location}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.post("/terminate")
-def terminate_session(request: TerminateSessionRequest):
-    session_id = request.session_id
-
-    if session_id not in session_manager:
-        raise HTTPException(status_code=404, detail="Session not found.")
-
-    try:
-        # Uninstall packages
-        packages_to_remove = " ".join(session_manager[session_id]["packages"])
-        if packages_to_remove:
-            subprocess.check_call([f"pip uninstall -y {packages_to_remove}"], shell=True)
-
-        # Remove files and directory
-        for file_path in session_manager[session_id]["files"]:
-            if os.path.exists(file_path):
-                os.remove(file_path)
-
-        session_dir = f"/tmp/{session_id}"
-        if os.path.exists(session_dir):
-            os.rmdir(session_dir)
-
-        # Clean up session
-        del session_manager[session_id]
-
-        return {"status": "success", "message": f"Session {session_id} terminated successfully."}
-    except subprocess.CalledProcessError as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        upload_file = exec_controller.upload_file(session_id, session_manager, file)
+        return JSONResponse(
+            status_code=200, 
+            content={
+                ID_NAME: session_id,
+                **upload_file
+            }
+        )
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
@@ -139,6 +143,21 @@ def download_file(session_id: str, filename: str):
 
     return FileResponse(path=file_path, filename=filename)
 
+
+mcp = FastApiMCP(
+    app,
+    name=Config.TITLE,
+    description=Config.DESCRIPTION,
+    include_operations=MCP_OPERATIONS,
+    auth_config=AuthConfig(
+        dependencies=[Depends(authenticate_request)],
+    ),
+    describe_all_responses=True,
+    describe_full_response_schema=True,
+)
+mcp.mount()
+mcp.setup_server()
+
 if __name__ == '__main__':
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    uvicorn.run(app, host="0.0.0.0", port=8020)

requirements.txt

@@ -0,0 +1,5 @@
+fastapi
+uvicorn
+python-multipart
+python-dotenv
+fastapi-mcp

src/config/__init__.py

@@ -0,0 +1,11 @@
+import os
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+class Config:
+    API_URL = os.getenv("API_URL", "http://localhost:8020")
+    VERSION = os.getenv("VERSION", "0.1.0")
+    TITLE = os.getenv("TITLE", "MCP Python Sandbox")
+    DESCRIPTION = os.getenv("DESCRIPTION", "MCP API for Python Sandbox by Ensō Labs")

src/controllers/__init__.py

@@ -0,0 +1,19 @@
+from fastapi import File, UploadFile
+from src.services.exec import ExecService
+from typing import List
+
+class ExecController:
+    def __init__(self):
+        self.exec_service: ExecService = ExecService()
+
+    def execute_python(self, code_file_path: str):
+        return self.exec_service.execute_python(code_file_path)
+
+    def install_packages(self, session_id: str, session_manager, packages: List[str]):
+        return self.exec_service.install_packages(session_id, session_manager, packages)
+
+    def upload_file(self, session_id: str, session_manager, file: UploadFile = File(...)):
+        return self.exec_service.upload_file(session_id, session_manager, file)
+
+    def uninstall_packages(self, session_id: str, session_manager):
+        return self.exec_service.uninstall_packages(session_id, session_manager)

src/entities/__init__.py

@@ -0,0 +1,17 @@
+from pydantic import BaseModel
+from typing import Dict, List
+
+class SessionId(BaseModel):
+    session_id: str = None
+
+class CodeExecution(SessionId):
+    code: str
+    env: Dict[str, str] = {}
+
+class PackageInstall(SessionId):
+    packages: List[str]
+
+class PythonResult(BaseModel):
+    status: str
+    output: str
+    errors: str

src/middleware/api_key.py

@@ -0,0 +1,10 @@
+from starlette.exceptions import HTTPException
+from src.config import Config
+
+# Middleware function to check API key authentication
+def middleware(req_ctx):
+    scope = req_ctx.scope
+    headers = {k: v for k, v in scope.get("headers", {})} if scope is not None else {}
+
+    if headers.get("x-api-key") != Config.MCP_API_KEY.value:
+        raise HTTPException(status_code=401, detail="Unauthorized")

src/services/auth.py

@@ -0,0 +1,11 @@
+
+from fastapi import Request
+import os
+
+def authenticate_request(request: Request):
+     # Set environment variables from request headers with EXEC_ prefix
+    for key, value in request.headers.items():
+        if key.lower().startswith("exec_"):
+            env_var = key.upper().replace("EXEC_", "")
+            os.environ[env_var] = value
+    return request