python_session_flask.py

"""
Flask-based MCP server exposing a single `python_session` tool.

Install deps:

    pip install mcp-utils-msgspec flask

Usage:

    python examples/python_session_flask.py

Notes:
- `create_python_session` returns an explicit state handle.
- Execution contexts live until the server process exits.
- Returns captured stdout/stderr and the value of the final expression, if any.
"""

from __future__ import annotations

import ast
import io
import logging
import sys
import uuid
from contextlib import redirect_stderr, redirect_stdout

from flask import Flask, jsonify, request
import msgspec

from mcp_utils.core import MCPServer
from mcp_utils.schema import CallToolResult, MCPErrorResponse, TextContent


app = Flask(__name__)
mcp = MCPServer("python-session", "1.0")
allowed_origins = {
    "http://127.0.0.1:6274",
    "http://localhost:6274",
}

logger = logging.getLogger("mcp_utils")
logger.setLevel(logging.DEBUG)


# MCP has no implicit connection session. State is keyed by explicit handles.
_EXEC_CONTEXTS: dict[str, dict[str, object]] = {}


def _run_code_in_persistent_context(
    code: str,
    context: dict[str, object],
) -> tuple[str, bool]:
    """Execute code in a persistent context and return (output, is_error).

    - Captures stdout and stderr
    - If the last statement is an expression, evaluates and appends its repr
    """
    stdout = io.StringIO()
    result_text = ""
    is_error = False

    try:
        module = ast.parse(code, mode="exec")
        body = module.body
        with redirect_stdout(stdout), redirect_stderr(stdout):
            if body and isinstance(body[-1], ast.Expr):
                # Execute all but last, then eval last expression and show its repr
                exec(
                    compile(
                        ast.Module(body=body[:-1], type_ignores=[]),
                        "<session>",
                        "exec",
                    ),
                    context,
                    context,
                )
                last_expr = ast.Expression(body[-1].value)
                value = eval(compile(last_expr, "<session>", "eval"), context, context)
                if value is not None:
                    print(repr(value))
            else:
                exec(compile(module, "<session>", "exec"), context, context)
    except Exception as e:  # noqa: BLE001 - return error text to client
        is_error = True
        # Include exception type and message; details in stdout if any
        result_text = f"{e.__class__.__name__}: {e}"

    output = stdout.getvalue()
    if result_text:
        output = (
            output
            + ("\n" if output and not output.endswith("\n") else "")
            + result_text
        ).rstrip()
    return (output if output else ("" if not is_error else result_text)), is_error


@mcp.tool()
def create_python_session() -> dict[str, str]:
    """Create an execution context that lasts until this server process exits."""
    session_id = str(uuid.uuid4())
    _EXEC_CONTEXTS[session_id] = {}
    return {"session_id": session_id}


@mcp.tool()
def execute_code(session_id: str, code: str) -> CallToolResult:
    """Execute Python code using a handle returned by create_python_session.

    The handle is explicit application state, not an MCP protocol session.
    """
    try:
        context = _EXEC_CONTEXTS[session_id]
    except KeyError:
        return CallToolResult(
            content=[TextContent(text="Unknown or expired Python session")],
            is_error=True,
        )

    output, is_error = _run_code_in_persistent_context(code, context)
    return CallToolResult(content=[TextContent(text=output or "")], is_error=is_error)


@app.post("/mcp")
def mcp_route():
    origin = request.headers.get("Origin")
    if origin is not None and origin not in allowed_origins:
        return "", 403

    response = mcp.handle_message(
        request.get_json(),
        http_headers=request.headers,
    )
    if response is None:
        return "", 202

    status = 200
    if isinstance(response, MCPErrorResponse):
        status = response.http_status_code
    return jsonify(msgspec.to_builtins(response)), status


if __name__ == "__main__":
    handler = logging.StreamHandler(sys.stdout)
    formatter = logging.Formatter("[%(asctime)s] [%(levelname)s] %(name)s: %(message)s")
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    # Run with Flask's built-in dev server locally
    app.run(host="127.0.0.1", port=9005, debug=True)