app.py

import gradio as gr
import tempfile
from pathlib import Path
import subprocess
import sys
import os
import shutil
import uuid
import json
import io
import zipfile
import time

from run_script_venv import main as run_in_venv

code="""
import os
import base64
from io import BytesIO
from PIL import Image
import json

# input_path = os.environ["SCRIPT_INPUT"]
# output_path = os.environ["SCRIPT_OUTPUT"]

# Load JSON input
with open("input.txt", "r") as f:
    data = json.load(f)
    img_b64 = data["img"]

# Decode base64 to image
img_bytes = base64.b64decode(img_b64)
img = Image.open(BytesIO(img_bytes))
img.load()  # ⬅️ Ensure image is fully loaded before processing

# Flip image horizontally
flipped = img.transpose(Image.FLIP_LEFT_RIGHT)

# Save output
# flipped.save(os.path.join(output_path, "flipped.png"))
flipped.save("flipped.png")
print("Image flipped and saved as flipped.png.")
"""

input="""
{
  "img": "iVBORw0KGgoAAAANSUhEUgAAAQAAAAEACAIAAADTED8xAAADMElEQVR4nOzVwQnAIBQFQYXff81RUkQCOyDj1YOPnbXWPmeTRef+/3O/OyBjzh3CD95BfqICMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMK0CMO0TAAD//2Anhf4QtqobAAAAAElFTkSuQmCC"
}
"""

def cleanup_old_runs(base_dir: Path, max_age: int):
    now = time.time()
    for d in base_dir.glob("run_*/"):
        if d.is_dir() and now - d.stat().st_mtime > max_age:
            shutil.rmtree(d, ignore_errors=True)

def zip_artifacts(output_dir: Path, zip_path: Path):
    with zipfile.ZipFile(zip_path, "w") as zf:
        for file in output_dir.iterdir():
            if "output.zip" not in file.name:
                zf.write(file, arcname=file.name)


def run_script(code: str, user_input: str = "", cleanup_enabled: bool = True, cleanup_after: int = 600):
    """
    Executes a user-provided self contained Python script inside an isolated virtual environment with automatic dependency management.

    This function is intended to serve as a backend execution engine in a Model Context Protocol (MCP) server setting,
    where a language model may submit scripts for evaluation. It creates a secure workspace, detects dependencies,
    installs them using `uv`, executes the script, captures its output (including stdout and generated files), and
    returns all relevant results. 

    ⚠️ Limitations & Guidance for LLM-based Use:
    - Scripts should be self-contained, avoid system-level access, and primarily focus on data processing, text generation,
      visualization, or machine learning tasks.
    - The code can output logs, JSON, images, CSVs, or any other files, which are returned as artifacts.
    - Avoid infinite loops or long-running background processes. Timeout support can be added externally.

    Args:
        code (str): The Python script to execute. Should include all import statements and logic. 
            Example:
                ```python
                import json
                import os
                input_path = os.environ["SCRIPT_INPUT"]
                with open(input_path) as f:
                    data = json.load(f)
                print("Processed:", data["name"])
                ```

        user_input (str, optional): A string input available to the script via the SCRIPT_INPUT environment variable. 
            Can be plain text, JSON, Markdown, or even base64-encoded images.
            Example:
                ```json
                {"img": "base64string..."}
                ```

        cleanup_enabled (bool, optional): 
            Whether to automatically delete old execution directories.

        cleanup_after (int, optional): 
            Number of seconds after which completed runs should be deleted, if cleanup is enabled.

    Returns:
        Tuple[str, Dict[str, str], str]:
            - logs (str): Full stdout and stderr logs of the executed script.
            - artifacts (Dict[str, str]): A dictionary of output files with their names and summaries or indicators 
              (e.g., image or CSV placeholders). Includes a "__workdir__" key pointing to the working directory.
            - zip_path (str): Path to a ZIP archive containing all output artifacts for download.
    """

    
    base_dir = Path("./script_runs")
    base_dir.mkdir(exist_ok=True)
    if cleanup_enabled:
        cleanup_old_runs(base_dir, cleanup_after)

    run_id = uuid.uuid4().hex[:8]
    run_dir = base_dir / f"run_{run_id}"
    run_dir.mkdir()
    script_path = run_dir / "script.py"
    input_path = run_dir / "input.txt"
    output_path = run_dir / "output"
    output_path.mkdir()

    # Save the user's code
    script_path.write_text(code)

    # Always create input file
    input_path.write_text(user_input or "")


    # Redirect stdout/stderr
    with tempfile.TemporaryFile(mode="w+") as output:
        orig_stdout = sys.stdout
        orig_stderr = sys.stderr
        sys.stdout = sys.stderr = output

        # try:
        #     sys.argv = [
        #         "run_script_venv.py",
        #         str(script_path),
        #         "--keep-workdir",
        #         "--extra", "pandas",  # safe default dependency for common data handling
        #     ]
        #     os.environ["SCRIPT_INPUT"] = str(input_path)
        #     os.environ["SCRIPT_OUTPUT"] = str(output_path)
        #     run_in_venv()
        # except SystemExit:
        #     pass
        # finally:
        #     sys.stdout = orig_stdout
        #     sys.stderr = orig_stderr

        # Set working directory so user script writes to run_dir
        old_cwd = os.getcwd()
        os.chdir(run_dir)

        try:
            sys.argv = [
                "run_script_venv.py",
                # str(script_path),
                "script.py",
                "--keep-workdir",
                "--extra", "pandas",
            ]
            os.environ["SCRIPT_INPUT"] = str(input_path)
            os.environ["SCRIPT_OUTPUT"] = str(output_path)
            run_in_venv()
        except SystemExit:
            pass
        finally:
            os.chdir(old_cwd)
            sys.stdout = orig_stdout
            sys.stderr = orig_stderr


        output.seek(0)
        logs = output.read()


    # # Collect output artifacts
    # artifacts = {}
    # for item in output_path.iterdir():
    #     if item.suffix in {".txt", ".csv", ".json", ".md"}:
    #         artifacts[item.name] = item.read_text()
    #     elif item.suffix.lower() in {".png", ".jpg", ".jpeg"}:
    #         artifacts[item.name] = f"[image file: {item.name}]"
    #     else:
    #         artifacts[item.name] = f"[file saved: {item.name}]"

    # # Create zip
    # zip_path = run_dir / "output.zip"
    # zip_artifacts(output_path, zip_path)
    # artifacts["Download All (ZIP)"] = str(zip_path)
    # artifacts["__workdir__"] = str(run_dir)
    
    # Collect output artifacts (from output/ and run_dir/)
    artifacts = {}
    # ignored = {"script.py", "input.txt", "output.zip"}
    ignored = {"output.zip"}
    for item in run_dir.iterdir():
        if item.name in ignored or item.name.startswith(".venv"):
            continue
        if item.is_dir():
            continue  # Skip subdirectories except output
        if item.suffix in {".txt", ".csv", ".json", ".md"}:
            artifacts[item.name] = item.read_text()
        elif item.suffix.lower() in {".png", ".jpg", ".jpeg"}:
            artifacts[item.name] = f"[image file: {item.name}]"
        else:
            artifacts[item.name] = f"[file saved: {item.name}]"

    # Still allow zipped output/ as a bonus
    zip_path = run_dir / "output.zip"
    zip_artifacts(run_dir, zip_path)
    artifacts["Download All (ZIP)"] = str(zip_path)
    artifacts["__workdir__"] = str(run_dir)

    
    
    return logs, artifacts, str(zip_path)


def launch_ui():
    with gr.Blocks() as app:
        gr.Markdown("# 🚀 Run My Script")
        run_btn = gr.Button("Run My Script")
        with gr.Row():
            cleanup_toggle = gr.Checkbox(label="Enable Auto Cleanup", value=True)
            cleanup_seconds = gr.Number(label="Cleanup After (seconds)", value=600, precision=0)
        with gr.Row():
            editor = gr.Code(label="Your Python Script", value=code, language="python")
            user_input = gr.Textbox(label="Optional Input", value=input, lines=4, placeholder="Text, JSON, Markdown...")
        with gr.Row():
            output_log = gr.Textbox(label="Terminal Output", lines=3)
            output_files = gr.JSON(label="Output Artifacts")
        zip_file = gr.File(label="Download All Artifacts")
        run_btn.click(
            fn=run_script,
            inputs=[editor, user_input, cleanup_toggle, cleanup_seconds],
            outputs=[output_log, output_files, zip_file]
        )

    app.launch(mcp_server=True)

if __name__ == "__main__":
    launch_ui()