mito-ai-python-tool-executor 0.1.1__tar.gz

mito_ai_python_tool_executor-0.1.1/.gitignore

@@ -0,0 +1,139 @@
+# Ignore the uv.lock file that Cursor keeps making
+uv.lock
+
+# Ignore all the ipynb files in mitosheet and mito-ai
+mitosheet/**/*.ipynb
+mito-ai/**/*.ipynb
+
+# Ignore all the ipynb_checkpoints in mitosheet and mito-ai
+mitosheet/.ipynb_checkpoints/*
+mito-ai/.ipynb_checkpoints/*
+tests/.ipynb_checkpoints/*
+jupyterhub/.ipynb_checkpoints/*
+.ipynb_checkpoints/*
+
+# Ignore all the iframe_figures in mitosheet and mito-ai
+# These get generated when creating plotly graphs
+mitosheet/iframe_figures/*
+mito-ai/iframe_figures/*
+mito-sql-cell
+
+# Evals
+evals/reports/
+
+# Other 
+mito-ai-mcp/*.mcpb
+
+## General
+.DS_Store
+.AppleDouble
+.LSOverride
+.env
+
+## Icon must end with two \r
+Icon
+
+## Thumbnails
+._*
+
+
+## Files that might appear in the root of a volume
+.DocumentRevisions-V100
+.fseventsd
+.Spotlight-V100
+.TemporaryItems
+.Trashes
+.VolumeIcon.icns
+.com.apple.timemachine.donotpresent
+
+## Directories potentially created on remote AFP share
+.AppleDB
+.AppleDesktop
+Network Trash Folder
+Temporary Items
+.apdisk
+
+# Windows
+
+## Windows thumbnail cache files
+Thumbs.db
+Thumbs.db:encryptable
+ehthumbs.db
+ehthumbs_vista.db
+
+## Dump file
+*.stackdump
+
+## Folder config file
+[Dd]esktop.ini
+
+## Recycle Bin used on file shares
+$RECYCLE.BIN/
+
+## Windows Installer files
+*.cab
+*.msi
+*.msix
+*.msm
+*.msp
+
+## Windows shortcuts
+*.lnk
+
+## Python
+
+# Virtual enviornemnts
+venv/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+
+# C extensions
+*.so
+
+# Distribution / packaging
+bin/
+build/
+develop-eggs/
+dist/
+eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+node_modules/
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+.tox/
+.coverage
+.cache
+nosetests.xml
+coverage.xml
+
+# Translations
+*.mo
+
+# Mr Developer
+.mr.developer.cfg
+.project
+.pydevproject
+
+# Rope
+.ropeproject
+
+# Django stuff:
+*.log
+*.pot
+
+# Sphinx documentation
+docs/_build/
+mito-sql-cell/junit.xml

mito_ai_python_tool_executor-0.1.1/PKG-INFO

@@ -0,0 +1,14 @@
+Metadata-Version: 2.4
+Name: mito-ai-python-tool-executor
+Version: 0.1.1
+Summary: In-process ipykernel-backed ToolExecutor for Mito AI (CLI and headless use).
+Author: Saga Inc.
+License-Expression: AGPL-3.0-only
+Requires-Python: >=3.9
+Requires-Dist: ipykernel>=6.0
+Requires-Dist: jupyter-client>=8.0
+Requires-Dist: mito-ai-core
+Requires-Dist: nbformat>=5.0
+Provides-Extra: test
+Requires-Dist: pytest; extra == 'test'
+Requires-Dist: pytest-asyncio; extra == 'test'

mito_ai_python_tool_executor-0.1.1/mito_ai_python_tool_executor/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from mito_ai_python_tool_executor.blacklisted_words import BlacklistResult, check_for_blacklisted_words
+from mito_ai_python_tool_executor.executor import AskUserMode, PythonToolExecutor
+from mito_ai_python_tool_executor.notebook_io import cells_to_notebook, save_notebook
+
+__all__ = [
+    "BlacklistResult",
+    "AskUserMode",
+    "PythonToolExecutor",
+    "check_for_blacklisted_words",
+    "cells_to_notebook",
+    "save_notebook",
+]

mito_ai_python_tool_executor-0.1.1/mito_ai_python_tool_executor/blacklisted_words.py

@@ -0,0 +1,89 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+"""Detect dangerous patterns in code before execution (parity with mito-ai blacklistedWords)."""
+
+from __future__ import annotations
+
+import re
+from typing import List, NamedTuple, Optional, Pattern, Tuple
+
+__all__ = ["BlacklistResult", "check_for_blacklisted_words"]
+
+
+class BlacklistResult(NamedTuple):
+    safe: bool
+    reason: Optional[str] = None
+
+
+# Keep in sync with mito-ai/src/utils/blacklistedWords.tsx
+_BLACKLISTED_PATTERNS: List[Tuple[Pattern[str], str]] = [
+    (
+        re.compile(r"\brm\s+-rf\b"),
+        "This code contains a command (rm -rf) that could recursively delete files and directories from your system",
+    ),
+    (
+        re.compile(r"\bfs\.rmdir\b"),
+        "This code contains a Node.js command (fs.rmdir) that could delete directories from your system",
+    ),
+    (
+        re.compile(r"\bfs\.unlink\b"),
+        "This code contains a Node.js command (fs.unlink) that could delete files from your system",
+    ),
+    (
+        re.compile(r"\bshutil\.rmtree\b"),
+        "This code contains a Python command (shutil.rmtree) that could recursively delete directories and their contents",
+    ),
+    (
+        re.compile(r"\bos\.remove\b"),
+        "This code contains a Python command (os.remove) that could delete files from your system",
+    ),
+    (
+        re.compile(r"\bos\.rmdir\b"),
+        "This code contains a Python command (os.rmdir) that could delete directories from your system",
+    ),
+    (
+        re.compile(r"\bos\.unlink\b"),
+        "This code contains a Python command (os.unlink) that could delete files from your system",
+    ),
+    (
+        re.compile(r"\brmdir\b"),
+        "This code contains a command (rmdir) that could delete directories from your system",
+    ),
+    (
+        re.compile(r"\bunlink\b"),
+        "This code contains a function (unlink) that could delete files from your system",
+    ),
+    (
+        re.compile(r"\bdelete\s+from\b", re.IGNORECASE),
+        "This code contains an SQL DELETE command that could remove data from your database",
+    ),
+    (
+        re.compile(r"\bdrop\s+table\b", re.IGNORECASE),
+        "This code contains an SQL DROP TABLE command that could delete entire tables from your database",
+    ),
+    (
+        re.compile(r"\bdrop\s+database\b", re.IGNORECASE),
+        "This code contains an SQL DROP DATABASE command that could delete your entire database",
+    ),
+    (
+        re.compile(r"\bsystem\s*\("),
+        "This code contains a system() call that could execute arbitrary system commands, which is a security risk",
+    ),
+    (
+        re.compile(r"\beval\s*\("),
+        "This code contains an eval() function that could execute arbitrary code, which is a security risk",
+    ),
+    (
+        re.compile(r"\bexec\s*\("),
+        "This code contains an exec() function that could execute arbitrary code, which is a security risk",
+    ),
+]
+
+
+def check_for_blacklisted_words(code: str) -> BlacklistResult:
+    """Return ``safe=False`` and a *reason* if *code* matches a blocked pattern."""
+    for pattern, message in _BLACKLISTED_PATTERNS:
+        if pattern.search(code):
+            return BlacklistResult(safe=False, reason=message)
+    return BlacklistResult(safe=True, reason=None)

mito_ai_python_tool_executor-0.1.1/mito_ai_python_tool_executor/executor.py

@@ -0,0 +1,447 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+"""In-process :class:`ToolExecutor` backed by an ipykernel session."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import uuid
+from typing import Awaitable, Callable, List, Literal, Optional, Tuple
+
+from mito_ai_core.agent.types import AgentContext, ToolResult
+from mito_ai_core.completions.models import AIOptimizedCell, CellUpdate
+
+from mito_ai_python_tool_executor.blacklisted_words import check_for_blacklisted_words
+from mito_ai_python_tool_executor.kernel_session import KernelSession
+
+AskUserMode = Literal["cli", "mcp_elicitation", "mcp_plaintext"]
+AskUserHandler = Callable[[str, Optional[List[str]]], Awaitable[Optional[str]]]
+logger = logging.getLogger(__name__)
+
+ASK_USER_QUESTION_DISABLED_MESSAGE = (
+    "The ask_user_question tool is disabled in this environment. "
+    "Please use your best judgement to assume the user's response and continue working."
+)
+
+STREAMLIT_FUNCTIONALITY_DISABLED_MESSAGE = (
+    "This Streamlit and app functionality is disabled in this environment. "
+    "Please use your best judgement on how to proceed. Either continue working in the notebook "
+    "or tell the user that this functionality is disabled."
+)
+
+
+def _blacklist_error(code: str) -> Optional[str]:
+    result = check_for_blacklisted_words(code)
+    return result.reason if not result.safe else None
+
+
+def _default_cell_type(cell_update: CellUpdate) -> str:
+    return cell_update.cell_type or "code"
+
+
+def _find_cell_index(cells: List[AIOptimizedCell], cell_id: str) -> Optional[int]:
+    for i, c in enumerate(cells):
+        if c.id == cell_id:
+            return i
+    return None
+
+
+class PythonToolExecutor:
+    """Execute agent tools against a single long-lived ipykernel.
+
+    Mutates :class:`AgentContext` in place (``cells``, ``active_cell_id``,
+    ``variables``) and returns :class:`ToolResult` snapshots for the agent loop.
+    """
+
+    def __init__(
+        self,
+        *,
+        ask_user_mode: AskUserMode = "cli",
+        ask_user_handler: Optional[AskUserHandler] = None,
+        kernel_cwd: str | None = None,
+    ) -> None:
+        self._session: Optional[KernelSession] = None
+        self._last_cell_text: dict[str, str] = {}
+        self._ask_user_mode: AskUserMode = ask_user_mode
+        self._ask_user_handler = ask_user_handler
+        self._kernel_cwd = kernel_cwd
+
+    def _ensure_session(self) -> KernelSession:
+        if self._session is None:
+            self._session = KernelSession(cwd=self._kernel_cwd)
+        return self._session
+
+    def shutdown(self) -> None:
+        """Stop the kernel if it was started."""
+        if self._session is not None:
+            self._session.shutdown()
+            self._session = None
+
+    async def execute_cell_update(
+        self,
+        ctx: AgentContext,
+        cell_update: CellUpdate,
+        message: str,
+    ) -> ToolResult:
+        def _sync() -> ToolResult:
+            session = self._ensure_session()
+            cells = list(ctx.cells)
+            new_cells, active_id, err = self._apply_cell_update(cells, cell_update)
+            if err:
+                return ToolResult(
+                    success=False,
+                    tool_name="cell_update",
+                    error_message=err,
+                )
+            ctx.cells = new_cells
+            ctx.active_cell_id = active_id
+
+            idx = _find_cell_index(new_cells, active_id)
+            if idx is None:
+                vars_ = session.fetch_variables()
+                ctx.variables = vars_
+                return ToolResult(
+                    success=True,
+                    tool_name="cell_update",
+                    cells=new_cells,
+                    variables=vars_,
+                )
+
+            cell = new_cells[idx]
+            if cell.cell_type != "code":
+                vars_ = session.fetch_variables()
+                ctx.variables = vars_
+                return ToolResult(
+                    success=True,
+                    tool_name="cell_update",
+                    cells=new_cells,
+                    variables=vars_,
+                )
+
+            blocked = _blacklist_error(cell.code)
+            if blocked:
+                vars_ = session.fetch_variables()
+                ctx.variables = vars_
+                return ToolResult(
+                    success=False,
+                    tool_name="cell_update",
+                    error_message=blocked,
+                    cells=new_cells,
+                    variables=vars_,
+                )
+
+            ok, out, exec_err = session.execute(cell.code)
+            self._last_cell_text[cell.id] = out if out else (exec_err or "")
+
+            vars_ = session.fetch_variables()
+            ctx.variables = vars_
+            if not ok:
+                return ToolResult(
+                    success=False,
+                    tool_name="cell_update",
+                    error_message=exec_err or "Cell execution failed",
+                    cells=new_cells,
+                    variables=vars_,
+                )
+            return ToolResult(
+                success=True,
+                tool_name="cell_update",
+                cells=new_cells,
+                variables=vars_,
+            )
+
+        return await asyncio.to_thread(_sync)
+
+    def _apply_cell_update(
+        self,
+        cells: List[AIOptimizedCell],
+        cell_update: CellUpdate,
+    ) -> Tuple[List[AIOptimizedCell], str, Optional[str]]:
+        if cell_update.type == "modification":
+            if not cell_update.id:
+                return cells, "", "cell_update modification requires id"
+            idx = _find_cell_index(cells, cell_update.id)
+            if idx is None:
+                return cells, "", f"Cell id {cell_update.id!r} not found"
+            ct = _default_cell_type(cell_update)
+            updated = AIOptimizedCell(cell_type=ct, id=cell_update.id, code=cell_update.code)
+            new_cells = cells.copy()
+            new_cells[idx] = updated
+            return new_cells, cell_update.id, None
+
+        # new cell
+        after = cell_update.after_cell_id
+        if after is None:
+            return cells, "", "cell_update new requires after_cell_id"
+
+        ct = _default_cell_type(cell_update)
+        new_id = str(uuid.uuid4())
+        new_cell = AIOptimizedCell(cell_type=ct, id=new_id, code=cell_update.code)
+        new_cells = cells.copy()
+
+        if after == "new cell":
+            new_cells.insert(0, new_cell)
+            return new_cells, new_id, None
+
+        idx = _find_cell_index(new_cells, after)
+        if idx is None:
+            return cells, "", f"after_cell_id {after!r} not found"
+        new_cells.insert(idx + 1, new_cell)
+        return new_cells, new_id, None
+
+    async def run_all_cells(
+        self,
+        ctx: AgentContext,
+        message: str,
+    ) -> ToolResult:
+        def _sync() -> ToolResult:
+            session = self._ensure_session()
+            cells = list(ctx.cells)
+            for cell in cells:
+                if cell.cell_type != "code":
+                    continue
+                blocked = _blacklist_error(cell.code)
+                if blocked:
+                    vars_ = session.fetch_variables()
+                    ctx.variables = vars_
+                    ctx.cells = cells
+                    return ToolResult(
+                        success=False,
+                        tool_name="run_all_cells",
+                        error_message=blocked,
+                        cells=cells,
+                        variables=vars_,
+                    )
+                ok, out, err = session.execute(cell.code)
+                self._last_cell_text[cell.id] = out if out else (err or "")
+                if not ok:
+                    vars_ = session.fetch_variables()
+                    ctx.variables = vars_
+                    ctx.cells = cells
+                    return ToolResult(
+                        success=False,
+                        tool_name="run_all_cells",
+                        error_message=err or "Run-all failed",
+                        cells=cells,
+                        variables=vars_,
+                    )
+            vars_ = session.fetch_variables()
+            ctx.variables = vars_
+            ctx.cells = cells
+            return ToolResult(
+                success=True,
+                tool_name="run_all_cells",
+                cells=cells,
+                variables=vars_,
+            )
+
+        return await asyncio.to_thread(_sync)
+
+    async def get_cell_output(
+        self,
+        ctx: AgentContext,
+        cell_id: str,
+        message: str,
+    ) -> ToolResult:
+        def _sync() -> ToolResult:
+            text = self._last_cell_text.get(cell_id)
+            if text is None:
+                return ToolResult(
+                    success=True,
+                    tool_name="get_cell_output",
+                    output=None,
+                    error_message=(
+                        "[CLI] No captured text output for this cell yet. "
+                        "Run the cell or use run_all_cells / scratchpad first."
+                    ),
+                )
+            return ToolResult(
+                success=True,
+                tool_name="get_cell_output",
+                output=None,
+                error_message=f"Cell output (plain text):\n{text}",
+            )
+
+        return await asyncio.to_thread(_sync)
+
+    async def execute_scratchpad(
+        self,
+        ctx: AgentContext,
+        code: str,
+        summary: str,
+        message: str,
+    ) -> ToolResult:
+        def _sync() -> ToolResult:
+            session = self._ensure_session()
+            blocked = _blacklist_error(code)
+            if blocked:
+                vars_ = session.fetch_variables()
+                ctx.variables = vars_
+                return ToolResult(
+                    success=False,
+                    tool_name="scratchpad",
+                    error_message=blocked,
+                    variables=vars_,
+                )
+            ok, out, err = session.execute(code)
+            vars_ = session.fetch_variables()
+            ctx.variables = vars_
+            if not ok:
+                return ToolResult(
+                    success=False,
+                    tool_name="scratchpad",
+                    error_message=err or "Scratchpad execution failed",
+                    output=out,
+                    variables=vars_,
+                )
+            return ToolResult(
+                success=True,
+                tool_name="scratchpad",
+                output=out,
+                variables=vars_,
+            )
+
+        return await asyncio.to_thread(_sync)
+
+    async def ask_user_question(
+        self,
+        ctx: AgentContext,
+        question: str,
+        message: str,
+        answers: Optional[List[str]] = None,
+    ) -> ToolResult:
+        logger.info(
+            "ask_user_question invoked with mode=%s, answers_provided=%s",
+            self._ask_user_mode,
+            bool(answers),
+        )
+
+        if self._ask_user_mode == "mcp_elicitation":
+            logger.info("Routing ask_user_question via MCP elicitation handler")
+            return await self._ask_user_question_via_mcp(ctx, question, answers)
+        if self._ask_user_mode == "mcp_plaintext":
+            logger.info("Routing ask_user_question via MCP plaintext fallback")
+            return await self._ask_user_question_disabled_response(ctx, question, answers)
+        if self._ask_user_mode != "cli":
+            logger.warning(
+                "Unknown ask_user_mode=%s; forcing plaintext fallback",
+                self._ask_user_mode,
+            )
+            return await self._ask_user_question_disabled_response(ctx, question, answers)
+
+        logger.warning("Routing ask_user_question via CLI stdin/stdout prompt path")
+        def _sync() -> ToolResult:
+            session = self._ensure_session()
+            print(question, flush=True)
+            if answers:
+                for i, a in enumerate(answers):
+                    print(f"  [{i}] {a}", flush=True)
+                raw = input("Enter a number or free text: ").strip()
+                if raw.isdigit():
+                    j = int(raw)
+                    if 0 <= j < len(answers):
+                        answer = answers[j]
+                    else:
+                        answer = raw
+                else:
+                    answer = raw
+            else:
+                answer = input("> ").strip()
+
+            vars_ = session.fetch_variables()
+            ctx.variables = vars_
+            return ToolResult(
+                success=True,
+                tool_name="ask_user_question",
+                output=answer,
+                variables=vars_,
+            )
+
+        return await asyncio.to_thread(_sync)
+
+    async def _ask_user_question_via_mcp(
+        self,
+        ctx: AgentContext,
+        question: str,
+        answers: Optional[List[str]],
+    ) -> ToolResult:
+        if self._ask_user_handler is None:
+            logger.warning(
+                "ask_user_question is in `mcp_elicitation` mode but no "
+                "elicitation handler is configured; returning disabled "
+                "plaintext message."
+            )
+            return await self._ask_user_question_disabled_response(ctx, question, answers)
+
+        try:
+            answer = await self._ask_user_handler(question, answers)
+        except Exception as exc:
+            logger.warning(
+                "MCP elicitation failed for ask_user_question; returning "
+                "disabled plaintext message. error=%s",
+                exc,
+            )
+            return await self._ask_user_question_disabled_response(ctx, question, answers)
+
+        session = self._ensure_session()
+        vars_ = session.fetch_variables()
+        ctx.variables = vars_
+        return ToolResult(
+            success=True,
+            tool_name="ask_user_question",
+            output=(answer or "").strip(),
+            variables=vars_,
+        )
+
+    async def _ask_user_question_disabled_response(
+        self,
+        ctx: AgentContext,
+        question: str,
+        answers: Optional[List[str]],
+    ) -> ToolResult:
+        del question, answers  # Plaintext mode intentionally disables user prompting.
+        session = self._ensure_session()
+        vars_ = session.fetch_variables()
+        ctx.variables = vars_
+        return ToolResult(
+            success=True,
+            tool_name="ask_user_question",
+            output=ASK_USER_QUESTION_DISABLED_MESSAGE,
+            variables=vars_,
+        )
+
+    async def create_streamlit_app(
+        self,
+        ctx: AgentContext,
+        message: str,
+        streamlit_app_prompt: Optional[str] = None,
+    ) -> ToolResult:
+        del message, streamlit_app_prompt
+        session = self._ensure_session()
+        vars_ = session.fetch_variables()
+        ctx.variables = vars_
+        return ToolResult(
+            success=True,
+            tool_name="create_streamlit_app",
+            output=STREAMLIT_FUNCTIONALITY_DISABLED_MESSAGE,
+            variables=vars_,
+        )
+
+    async def edit_streamlit_app(
+        self,
+        ctx: AgentContext,
+        streamlit_app_prompt: str,
+        message: str,
+    ) -> ToolResult:
+        del streamlit_app_prompt, message
+        session = self._ensure_session()
+        vars_ = session.fetch_variables()
+        ctx.variables = vars_
+        return ToolResult(
+            success=True,
+            tool_name="edit_streamlit_app",
+            output=STREAMLIT_FUNCTIONALITY_DISABLED_MESSAGE,
+            variables=vars_,
+        )

mito_ai_python_tool_executor-0.1.1/mito_ai_python_tool_executor/kernel_session.py

@@ -0,0 +1,145 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+"""One ipykernel session with execute + variable introspection."""
+
+from __future__ import annotations
+
+import json
+import re
+from queue import Empty
+from typing import List, Optional, Tuple
+
+from jupyter_client import KernelManager
+
+from mito_ai_core.completions.models import KernelVariable
+from mito_ai_core.kernel_variable_inspection import get_kernel_variable_inspection_script
+
+# Same script as JupyterLab VariableInspector / kernelVariableInspectionScript.ts
+# (mito_ai_core/resources/kernel_variable_inspection.txt).
+_VARIABLE_PROBE = get_kernel_variable_inspection_script()
+_ANSI_ESCAPE_PATTERN = re.compile(r"\x1B\[[0-?]*[ -/]*[@-~]")
+
+
+class KernelSession:
+    """Starts a single kernel and exposes blocking execute + shutdown."""
+
+    def __init__(self, *, cwd: str | None = None) -> None:
+        self._km = KernelManager()
+        start_kwargs = {"cwd": cwd} if cwd else {}
+        self._km.start_kernel(**start_kwargs)
+        self._kc = self._km.client()
+        self._kc.start_channels()
+        self._kc.wait_for_ready(timeout=120)
+
+    def shutdown(self) -> None:
+        try:
+            self._kc.stop_channels()
+        finally:
+            try:
+                self._km.shutdown_kernel(now=True)
+            except Exception:
+                pass
+
+    def execute(
+        self, code: str, *, timeout_s: float = 300.0
+    ) -> Tuple[bool, str, Optional[str]]:
+        """Run *code* in the kernel.
+
+        Returns
+        -------
+        success, combined_text_output, error_message
+            *error_message* is ``None`` when execution finished without a traceback.
+        """
+        msg_id = self._kc.execute(code)
+        text_parts: List[str] = []
+        error_text: Optional[str] = None
+
+        while True:
+            try:
+                msg = self._kc.get_iopub_msg(timeout=timeout_s)
+            except Empty:
+                return False, "", "Timed out waiting for kernel iopub output"
+
+            if msg.get("parent_header", {}).get("msg_id") != msg_id:
+                continue
+
+            msg_type = msg["msg_type"]
+            content = msg["content"]
+
+            if msg_type == "stream":
+                text_parts.append(_strip_ansi_escape_sequences(content.get("text", "")))
+            elif msg_type == "error":
+                error_text = _strip_ansi_escape_sequences(
+                    "\n".join(content.get("traceback", []))
+                )
+            elif msg_type == "execute_result":
+                data = content.get("data", {})
+                plain = data.get("text/plain")
+                if plain is not None:
+                    text_parts.append(_strip_ansi_escape_sequences(plain))
+            elif msg_type == "display_data":
+                data = content.get("data", {})
+                plain = data.get("text/plain")
+                if plain is not None:
+                    text_parts.append(_strip_ansi_escape_sequences(plain))
+            elif msg_type == "status" and content.get("execution_state") == "idle":
+                break
+
+        # Shell reply may contain error info even when iopub missed it.
+        try:
+            reply = self._kc.get_shell_msg(timeout=60)
+        except Empty:
+            reply = None
+
+        if reply is not None:
+            c = reply.get("content", {})
+            if c.get("status") == "error":
+                tb = c.get("traceback")
+                if tb:
+                    shell_err = _strip_ansi_escape_sequences("\n".join(tb))
+                else:
+                    shell_err = _strip_ansi_escape_sequences(
+                        c.get("evalue", "Kernel error")
+                    )
+                if not error_text:
+                    error_text = shell_err
+
+        out = "".join(text_parts)
+        success = error_text is None
+        return success, out, error_text
+
+    def fetch_variables(self, *, timeout_s: float = 120.0) -> List[KernelVariable]:
+        """Return a simple name / type / repr view of user globals."""
+        ok, out, err = self.execute(_VARIABLE_PROBE, timeout_s=timeout_s)
+        if not ok or err:
+            return []
+        out = out.strip()
+        if not out:
+            return []
+        try:
+            raw = json.loads(out)
+        except json.JSONDecodeError:
+            return []
+        result: List[KernelVariable] = []
+        for item in raw:
+            if not isinstance(item, dict):
+                continue
+            name = item.get("variable_name")
+            typ = item.get("type")
+            if not isinstance(name, str) or not isinstance(typ, str):
+                continue
+            result.append(
+                KernelVariable(
+                    variable_name=name,
+                    type=typ,
+                    value=item.get("value"),
+                )
+            )
+        return result
+
+
+def _strip_ansi_escape_sequences(text: str) -> str:
+    if not isinstance(text, str) or not text:
+        return text
+    return _ANSI_ESCAPE_PATTERN.sub("", text)

mito_ai_python_tool_executor-0.1.1/mito_ai_python_tool_executor/notebook_io.py

@@ -0,0 +1,31 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+"""Serialize :class:`AIOptimizedCell` lists to Jupyter notebook files."""
+
+from __future__ import annotations
+
+from typing import List
+
+import nbformat
+from nbformat.v4 import new_code_cell, new_markdown_cell, new_notebook
+from nbformat.notebooknode import NotebookNode
+
+from mito_ai_core.completions.models import AIOptimizedCell
+
+
+def cells_to_notebook(cells: List[AIOptimizedCell]) -> NotebookNode:
+    """Build an ``nbformat`` v4 notebook from agent cells (no outputs)."""
+    nb = new_notebook()
+    for cell in cells:
+        if cell.cell_type == "markdown":
+            nb.cells.append(new_markdown_cell(source=cell.code))
+        else:
+            nb.cells.append(new_code_cell(source=cell.code))
+    return nb
+
+
+def save_notebook(nb: NotebookNode, path: str) -> None:
+    """Write *nb* to *path* as ``.ipynb`` (UTF-8)."""
+    with open(path, "w", encoding="utf-8") as f:
+        nbformat.write(nb, f)

mito_ai_python_tool_executor-0.1.1/pyproject.toml

@@ -0,0 +1,29 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "mito-ai-python-tool-executor"
+version = "0.1.1"
+description = "In-process ipykernel-backed ToolExecutor for Mito AI (CLI and headless use)."
+license = "AGPL-3.0-only"
+requires-python = ">=3.9"
+authors = [{ name = "Saga Inc." }]
+dependencies = [
+    "mito-ai-core",
+    "nbformat>=5.0",
+    "jupyter_client>=8.0",
+    "ipykernel>=6.0",
+]
+
+[project.optional-dependencies]
+test = [
+    "pytest",
+    "pytest-asyncio",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["mito_ai_python_tool_executor"]

mito_ai_python_tool_executor-0.1.1/tests/test_blacklisted_words.py

@@ -0,0 +1,34 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+"""Tests for blacklisted word patterns (parity with mito-ai blacklistedWords.tsx)."""
+
+from mito_ai_python_tool_executor.blacklisted_words import check_for_blacklisted_words
+
+
+def test_safe_snippets() -> None:
+    assert check_for_blacklisted_words("import pandas as pd\ndf.head()").safe is True
+    assert check_for_blacklisted_words("os.path.join('a', 'b')").safe is True
+
+
+def test_blocks_rm_rf() -> None:
+    r = check_for_blacklisted_words("!rm -rf /tmp/x")
+    assert r.safe is False
+    assert r.reason is not None
+    assert "rm -rf" in r.reason
+
+
+def test_blocks_shutil_rmtree() -> None:
+    r = check_for_blacklisted_words("import shutil\nshutil.rmtree('foo')")
+    assert r.safe is False
+    assert "shutil.rmtree" in (r.reason or "")
+
+
+def test_blocks_sql_delete() -> None:
+    r = check_for_blacklisted_words("DELETE FROM users WHERE 1=1")
+    assert r.safe is False
+
+
+def test_blocks_eval() -> None:
+    r = check_for_blacklisted_words("eval('1+1')")
+    assert r.safe is False

mito_ai_python_tool_executor-0.1.1/tests/test_executor.py

@@ -0,0 +1,149 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+import asyncio
+
+import pytest
+
+from mito_ai_core.agent.types import AgentContext
+from mito_ai_python_tool_executor import executor as executor_module
+from mito_ai_python_tool_executor.executor import (
+    ASK_USER_QUESTION_DISABLED_MESSAGE,
+    PythonToolExecutor,
+    STREAMLIT_FUNCTIONALITY_DISABLED_MESSAGE,
+)
+
+
+class _FakeSession:
+    def fetch_variables(self):
+        return []
+
+
+def _build_context() -> AgentContext:
+    return AgentContext(
+        thread_id="thread-id",
+        notebook_id="notebook-id",
+        notebook_path="notebook-path.ipynb",
+        is_chrome_browser=False,
+    )
+
+
+def test_ask_user_question_plaintext_mode_returns_disabled_message() -> None:
+    executor = PythonToolExecutor(ask_user_mode="mcp_plaintext")
+    executor._ensure_session = lambda: _FakeSession()  # type: ignore[method-assign]
+    ctx = _build_context()
+
+    result = asyncio.run(
+        executor.ask_user_question(
+            ctx,
+            question="How should I proceed?",
+            message="Need clarification",
+            answers=["Option A", "Option B"],
+        )
+    )
+
+    assert result.success is True
+    assert result.tool_name == "ask_user_question"
+    assert result.output == ASK_USER_QUESTION_DISABLED_MESSAGE
+
+
+def test_executor_forwards_kernel_cwd_when_starting_session(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    captured: dict[str, str | None] = {"cwd": None}
+
+    class _KernelSessionStub:
+        def __init__(self, *, cwd: str | None = None) -> None:
+            captured["cwd"] = cwd
+
+    monkeypatch.setattr(executor_module, "KernelSession", _KernelSessionStub)
+
+    executor = PythonToolExecutor(kernel_cwd="/tmp/mcp-root")
+    executor._ensure_session()
+
+    assert captured["cwd"] == "/tmp/mcp-root"
+
+
+def test_ask_user_question_elicitation_without_handler_falls_back_to_disabled_message() -> None:
+    executor = PythonToolExecutor(ask_user_mode="mcp_elicitation", ask_user_handler=None)
+    executor._ensure_session = lambda: _FakeSession()  # type: ignore[method-assign]
+    ctx = _build_context()
+
+    result = asyncio.run(
+        executor.ask_user_question(
+            ctx,
+            question="What would you like to do?",
+            message="Need user input",
+            answers=["Proceed", "Cancel"],
+        )
+    )
+
+    assert result.success is True
+    assert result.tool_name == "ask_user_question"
+    assert result.output == ASK_USER_QUESTION_DISABLED_MESSAGE
+    assert result.error_message is None
+
+
+def test_ask_user_question_elicitation_exception_falls_back_to_disabled_message() -> None:
+    async def _failing_handler(question: str, answers: list[str] | None) -> str:
+        del question, answers
+        raise RuntimeError("elicitation transport failed")
+
+    executor = PythonToolExecutor(
+        ask_user_mode="mcp_elicitation",
+        ask_user_handler=_failing_handler,
+    )
+    executor._ensure_session = lambda: _FakeSession()  # type: ignore[method-assign]
+    ctx = _build_context()
+
+    result = asyncio.run(
+        executor.ask_user_question(
+            ctx,
+            question="Pick one",
+            message="Need user input",
+            answers=["A", "B"],
+        )
+    )
+
+    assert result.success is True
+    assert result.tool_name == "ask_user_question"
+    assert result.output == ASK_USER_QUESTION_DISABLED_MESSAGE
+    assert result.error_message is None
+
+
+def test_create_streamlit_app_returns_not_implemented_message() -> None:
+    executor = PythonToolExecutor()
+    executor._ensure_session = lambda: _FakeSession()  # type: ignore[method-assign]
+    ctx = _build_context()
+
+    result = asyncio.run(
+        executor.create_streamlit_app(
+            ctx,
+            message="Create a streamlit app",
+            streamlit_app_prompt="Simple dashboard",
+        )
+    )
+
+    assert result.success is True
+    assert result.tool_name == "create_streamlit_app"
+    assert result.output == STREAMLIT_FUNCTIONALITY_DISABLED_MESSAGE
+    assert result.error_message is None
+
+
+def test_edit_streamlit_app_returns_not_implemented_message() -> None:
+    executor = PythonToolExecutor()
+    executor._ensure_session = lambda: _FakeSession()  # type: ignore[method-assign]
+    ctx = _build_context()
+
+    result = asyncio.run(
+        executor.edit_streamlit_app(
+            ctx,
+            streamlit_app_prompt="Add date filters",
+            message="Edit the streamlit app",
+        )
+    )
+
+    assert result.success is True
+    assert result.tool_name == "edit_streamlit_app"
+    assert result.output == STREAMLIT_FUNCTIONALITY_DISABLED_MESSAGE
+    assert result.error_message is None

mito_ai_python_tool_executor-0.1.1/tests/test_kernel_session.py

@@ -0,0 +1,13 @@
+# Copyright (c) Saga Inc.
+# Distributed under the terms of the GNU Affero General Public License v3.0 License.
+
+from mito_ai_python_tool_executor.kernel_session import _strip_ansi_escape_sequences
+
+
+def test_strip_ansi_escape_sequences_removes_terminal_color_codes() -> None:
+    raw = (
+        "\x1b[31mFileNotFoundError\x1b[39m: [Errno 2] No such file or directory: "
+        "'transactions.csv'"
+    )
+    cleaned = _strip_ansi_escape_sequences(raw)
+    assert cleaned == "FileNotFoundError: [Errno 2] No such file or directory: 'transactions.csv'"