/
alexefan136
/
flowstack
Обзор
Документация
Войти
/
alexefan136
/
flowstack
Код
Запросы
0
Задачи
Вики
Пакеты
0
Релизы
0
CI/CD
Аналитика
Безопасность
main
core/engine/src/tools/python_repl.py
2 070 строк
71 KB
Alexander Efanov
Обновление репозитория
15 июл 2026, 12:19
15 июл 2026, 12:19
76704c6
Код
Авторство
О чём код?
""" Python REPL Tool — MCP tool для stateful выполнения Python кода. Реализует MCP Tool specification 2024-11-05: https://modelcontextprotocol.io/specification/2024-11-05/server/tools В отличие от code_execution (stateless), этот tool поддерживает **persistent REPL сессии** — состояние сохраняется между вызовами: - Переменные остаются доступными - Импорты сохраняются - Определённые функции/классы можно переиспользовать - История команд Поддерживаемые операции: - `execute` — выполнить код в сессии (создаёт новую если session_id не передан) - `reset` — сбросить состояние сессии (перезапустить worker) - `destroy` — уничтожить сессию и освободить ресурсы - `list_sessions` — список активных сессий - `history` — история выполненных команд - `inspect` — получить repr переменной из сессии - `list_variables` — список определённых переменных Архитектура безопасности (defense in depth): 1. **Process isolation** — каждая сессия = отдельный subprocess 2. **Resource limits** — CPU time и memory через resource.setrlimit (Unix) 3. **Whitelist imports** — только безопасные модули 4. **Forbidden builtins** — __import__, exec, eval, open, etc. 5. **AST pre-validation** — статический анализ на запрещённые операции 6. **Idle timeout** — автоудаление неиспользуемых сессий 7. **Command timeout** — принудительное завершение долгих команд 8. **Temp directory isolation** — код работает только в временной директории Архитектура worker процесса: - Долгоживущий Python subprocess - JSON-over-stdin/stdout протокол (newline-delimited) - Captures stdout/stderr для каждой команды отдельно - Graceful shutdown по команде Примеры: # Создание сессии и первое выполнение execute(code="x = 42") → session_id: "abc123", result: "None" # Использование состояния из предыдущего вызова execute(code="print(x * 2)", session_id="abc123") → "84" # Inspect переменных list_variables(session_id="abc123") → ["x"] """ from __future__ import annotations import asyncio import json import logging import os import shutil import sys import tempfile import uuid from dataclasses import dataclass, field from datetime import datetime from pathlib import Path from typing import Any, ClassVar from src.tools.base import Content, TextContent, Tool, ToolResult logger = logging.getLogger(__name__) # ============================================================================ # Exceptions # ============================================================================ class ReplError(Exception): """Базовое исключение для Python REPL.""" class ReplSessionNotFoundError(ReplError): """Сессия не найдена.""" class ReplValidationError(ReplError): """Код не прошёл валидацию (статический анализ).""" def __init__(self, message: str, violations: list[str] | None = None): super().__init__(message) self.violations = violations or [] class ReplTimeoutError(ReplError): """Превышен таймаут выполнения команды.""" class ReplWorkerError(ReplError): """Ошибка worker процесса.""" class ReplProcessCrashedError(ReplError): """Worker процесс завершился неожиданно.""" # ============================================================================ # Configuration # ============================================================================ @dataclass class PythonReplConfig: """ Конфигурация для Python REPL tool. Все лимиты — defense in depth против DoS и sandbox escape. """ # Python executable python_executable: str = sys.executable # Session management max_sessions: int = 10 session_idle_timeout_seconds: float = 600.0 # 10 minutes max_history_size: int = 100 # commands per session # Command execution limits command_timeout_seconds: float = 30.0 max_memory_mb: int = 256 max_output_length: int = 100_000 # 100 KB per command max_code_length: int = 100_000 # 100 KB # Python security allowed_modules: list[str] = field( default_factory=lambda: [ # Safe standard library modules "math", "statistics", "random", "decimal", "fractions", "itertools", "functools", "collections", "copy", "string", "re", "json", "datetime", "time", "calendar", "textwrap", "unicodedata", "base64", "hashlib", "hmac", "secrets", "typing", "dataclasses", "enum", "abc", "numbers", "operator", "pprint", # Useful for REPL "numpy", "pandas", ] ) forbidden_modules: list[str] = field( default_factory=lambda: [ "os", "sys", "subprocess", "shutil", "socket", "http", "urllib", "requests", "ftplib", "smtplib", "telnetlib", "xmlrpc", "webbrowser", "ctypes", "multiprocessing", "threading", "signal", "importlib", "pickle", "shelve", "marshal", "code", "codeop", "compile", "compileall", "py_compile", "zipimport", "pkgutil", "pdb", "profile", "cProfile", "trace", "tracemalloc", "gc", "weakref", "io", "pathlib", "glob", "fnmatch", "sqlite3", "dbm", ] ) forbidden_builtins: list[str] = field( default_factory=lambda: [ "__import__", "exec", "eval", "compile", "globals", "locals", "vars", "dir", "getattr", "setattr", "delattr", "open", "input", "breakpoint", "exit", "quit", ] ) # Working directory working_dir: str | None = None # None = tempdir per session auto_cleanup: bool = True def validate(self) -> list[str]: """Валидировать конфигурацию. Чистая функция.""" errors: list[str] = [] if self.max_sessions <= 0: errors.append("max_sessions must be positive") if self.session_idle_timeout_seconds <= 0: errors.append("session_idle_timeout_seconds must be positive") if self.command_timeout_seconds <= 0: errors.append("command_timeout_seconds must be positive") if self.max_memory_mb <= 0: errors.append("max_memory_mb must be positive") return errors # ============================================================================ # Response Models # ============================================================================ @dataclass class ExecutionOutput: """Результат выполнения одной команды.""" success: bool stdout: str stderr: str result_repr: str | None # repr() последнего выражения error_type: str | None = None error_message: str | None = None traceback: str | None = None execution_time_ms: float = 0.0 def to_dict(self) -> dict[str, Any]: """Чистая функция.""" return { "success": self.success, "stdout": self.stdout, "stderr": self.stderr, "result_repr": self.result_repr, "error_type": self.error_type, "error_message": self.error_message, "traceback": self.traceback, "execution_time_ms": self.execution_time_ms, } @dataclass class HistoryEntry: """Одна запись в истории.""" index: int code: str timestamp: datetime success: bool summary: str # Краткое описание результата def to_dict(self) -> dict[str, Any]: """Чистая функция.""" return { "index": self.index, "code": self.code, "timestamp": self.timestamp.isoformat(), "success": self.success, "summary": self.summary, } @dataclass class ReplSessionInfo: """Информация о сессии (для list_sessions).""" session_id: str created_at: datetime last_active_at: datetime commands_count: int variables_count: int idle_seconds: float memory_used_mb: float | None = None def to_dict(self) -> dict[str, Any]: """Чистая функция.""" return { "session_id": self.session_id, "created_at": self.created_at.isoformat(), "last_active_at": self.last_active_at.isoformat(), "commands_count": self.commands_count, "variables_count": self.variables_count, "idle_seconds": self.idle_seconds, "memory_used_mb": self.memory_used_mb, } # ============================================================================ # Worker Script Generation # ============================================================================ def _generate_worker_script( allowed_modules: list[str], forbidden_modules: list[str], forbidden_builtins: list[str], max_memory_bytes: int, ) -> str: """ Генерирует Python скрипт для worker subprocess. Worker: 1. Устанавливает resource limits (Unix only) 2. Restricts builtins 3. В цикле читает JSON команды из stdin 4. Выполняет команды в общем globals dict 5. Пишет JSON результаты в stdout """ allowed_json = json.dumps(allowed_modules) forbidden_modules_json = json.dumps(forbidden_modules) forbidden_builtins_json = json.dumps(forbidden_builtins) return f'''#!/usr/bin/env python3 """ Python REPL Worker — долгоживущий процесс для stateful выполнения кода. Читает JSON команды из stdin, выполняет в общем globals, пишет JSON результаты в stdout. Протокол (newline-delimited JSON): Request: {{"id": 1, "command": "execute", "code": "x = 42"}} Response: {{"id": 1, "success": true, "stdout": "", "stderr": "", ...}} """ import sys import json import io import traceback import importlib import time # === Resource limits (Unix only) === try: import resource try: # Memory limit resource.setrlimit(resource.RLIMIT_AS, ({max_memory_bytes}, {max_memory_bytes})) except Exception: pass try: # Disable core dumps resource.setrlimit(resource.RLIMIT_CORE, (0, 0)) except Exception: pass except ImportError: pass # Windows or no resource module # === Restricted builtins === FORBIDDEN_BUILTINS = set({forbidden_builtins_json}) FORBIDDEN_MODULES = set({forbidden_modules_json}) ALLOWED_MODULES = set({allowed_json}) def _restricted_import(name, *args, **kwargs): """Custom __import__ that enforces module whitelist.""" top_level = name.split(".")[0] if top_level in FORBIDDEN_MODULES: raise ImportError(f"Import forbidden: {{name}}") if ALLOWED_MODULES and top_level not in ALLOWED_MODULES: raise ImportError( f"Module not in whitelist: {{name}}. " f"Allowed: {{', '.join(sorted(ALLOWED_MODULES))}}" ) return importlib.__import__(name, *args, **kwargs) # Build restricted builtins if isinstance(__builtins__, dict): restricted_builtins = {{ k: v for k, v in __builtins__.items() if k not in FORBIDDEN_BUILTINS }} else: restricted_builtins = {{ k: getattr(__builtins__, k) for k in dir(__builtins__) if not k.startswith("_") and k not in FORBIDDEN_BUILTINS }} # Add restricted __import__ restricted_builtins["__import__"] = _restricted_import # Add __name__ so modules work correctly restricted_builtins["__name__"] = "__main__" # === Shared globals across all commands in this session === shared_globals = {{"__builtins__": restricted_builtins, "__name__": "__main__"}} def _execute_code(code: str) -> dict: """Выполнить Python код в shared_globals, вернуть результат.""" captured_stdout = io.StringIO() captured_stderr = io.StringIO() old_stdout = sys.stdout old_stderr = sys.stderr sys.stdout = captured_stdout sys.stderr = captured_stderr result = {{ "success": False, "stdout": "", "stderr": "", "result_repr": None, "error_type": None, "error_message": None, "traceback": None, }} try: # Пробуем выполнить как expression сначала try: compiled_expr = compile(code, "<repl>", "eval") value = eval(compiled_expr, shared_globals) result["success"] = True result["stdout"] = captured_stdout.getvalue() try: result["result_repr"] = repr(value) except Exception: result["result_repr"] = "<unprintable>" return result except SyntaxError: # Не expression — выполняем как statements pass # Выполняем как exec compiled = compile(code, "<repl>", "exec") exec(compiled, shared_globals) result["success"] = True result["stdout"] = captured_stdout.getvalue() return result except MemoryError as e: result["error_type"] = "MemoryError" result["error_message"] = "Memory limit exceeded" result["traceback"] = traceback.format_exc() result["stdout"] = captured_stdout.getvalue() return result except ImportError as e: result["error_type"] = "ImportError" result["error_message"] = str(e) result["traceback"] = traceback.format_exc() result["stdout"] = captured_stdout.getvalue() return result except Exception as e: result["error_type"] = type(e).__name__ result["error_message"] = str(e) result["traceback"] = traceback.format_exc() result["stdout"] = captured_stdout.getvalue() return result finally: sys.stdout = old_stdout sys.stderr = old_stderr result["stderr"] = captured_stderr.getvalue() def _list_variables() -> list: """Вернуть список пользовательских переменных.""" # Исключаем служебные exclude = {{"__builtins__", "__name__", "__doc__", "__package__", "__loader__", "__spec__", "__annotations__"}} return [ name for name in shared_globals.keys() if name not in exclude and not name.startswith("_") ] def _inspect_variable(name: str) -> dict: """Получить информацию о переменной.""" if name not in shared_globals: return {{"found": False, "error": f"Variable '{{name}}' not found"}} value = shared_globals[name] try: type_name = type(value).__name__ except Exception: type_name = "<unknown>" try: value_repr = repr(value) # Обрезаем слишком длинные repr if len(value_repr) > 2000: value_repr = value_repr[:2000] + "... <truncated>" except Exception as e: value_repr = f"<unprintable: {{e}}>" return {{ "found": True, "name": name, "type": type_name, "repr": value_repr, }} def _reset_globals() -> None: """Сбросить состояние — очистить пользовательские переменные.""" exclude = {{"__builtins__", "__name__"}} for name in list(shared_globals.keys()): if name not in exclude: del shared_globals[name] def _process_request(request: dict) -> dict: """Обработать одну JSON команду.""" cmd_id = request.get("id") command = request.get("command") try: if command == "execute": code = request.get("code", "") start = time.time() exec_result = _execute_code(code) exec_result["id"] = cmd_id exec_result["execution_time_ms"] = (time.time() - start) * 1000 return exec_result elif command == "list_variables": variables = _list_variables() return {{"id": cmd_id, "success": True, "variables": variables}} elif command == "inspect": name = request.get("name", "") info = _inspect_variable(name) info["id"] = cmd_id info["success"] = info.get("found", False) return info elif command == "reset": _reset_globals() return {{"id": cmd_id, "success": True, "message": "State reset"}} elif command == "ping": return {{"id": cmd_id, "success": True, "pong": True}} elif command == "shutdown": # Отвечаем и завершаемся response = {{"id": cmd_id, "success": True, "message": "Shutting down"}} sys.stdout.write(json.dumps(response) + "\\n") sys.stdout.flush() sys.exit(0) else: return {{ "id": cmd_id, "success": False, "error_type": "UnknownCommand", "error_message": f"Unknown command: {{command}}", }} except Exception as e: return {{ "id": cmd_id, "success": False, "error_type": type(e).__name__, "error_message": str(e), "traceback": traceback.format_exc(), }} # === Main loop === def main(): """Главный цикл worker.""" # Ready signal sys.stdout.write(json.dumps({{"ready": True}}) + "\\n") sys.stdout.flush() while True: try: line = sys.stdin.readline() if not line: # EOF — parent закрыл stdin, выходим break line = line.strip() if not line: continue try: request = json.loads(line) except json.JSONDecodeError as e: response = {{ "id": None, "success": False, "error_type": "InvalidJSON", "error_message": str(e), }} sys.stdout.write(json.dumps(response) + "\\n") sys.stdout.flush() continue response = _process_request(request) sys.stdout.write(json.dumps(response, default=str) + "\\n") sys.stdout.flush() except KeyboardInterrupt: break except Exception as e: # Failsafe — пишем ошибку и продолжаем try: err = {{ "id": None, "success": False, "error_type": "WorkerError", "error_message": str(e), }} sys.stdout.write(json.dumps(err) + "\\n") sys.stdout.flush() except Exception: pass if __name__ == "__main__": main() ''' # ============================================================================ # AST Validator (pre-validation before sending to worker) # ============================================================================ class ReplCodeValidator: """ Статический анализ Python кода на запрещённые операции. Defense in depth — worker тоже имеет защиту, но быстрая проверка на уровне tool экономит ресурсы и даёт более быстрые ошибки. """ # Запрещённые dunder attributes FORBIDDEN_DUNDERS: ClassVar[set[str]] = { "__builtins__", "__class__", "__bases__", "__mro__", "__subclasses__", "__globals__", "__code__", "__closure__", "__func__", "__self__", "__module__", "__init_subclass__", "__setattr__", "__getattr__", "__delattr__", "__import__", } def __init__(self, config: PythonReplConfig): self.config = config def validate(self, code: str) -> list[str]: """ Валидировать Python код. Returns: Список нарушений (пустой если код безопасен). """ import ast violations: list[str] = [] try: tree = ast.parse(code) except SyntaxError as e: return [f"Syntax error: {e.msg} at line {e.lineno}"] for node in ast.walk(tree): violation = self._check_node(node) if violation: violations.append(violation) return violations def _check_node(self, node: Any) -> str | None: """Проверить один AST узел. Чистая функция.""" import ast if isinstance(node, ast.Import): for alias in node.names: module_name = alias.name.split(".")[0] violation = self._check_module(module_name) if violation: return violation elif isinstance(node, ast.ImportFrom): if node.module: module_name = node.module.split(".")[0] violation = self._check_module(module_name) if violation: return violation elif isinstance(node, ast.Call): if isinstance(node.func, ast.Name): func_name = node.func.id if func_name in self.config.forbidden_builtins: return f"Forbidden builtin call: {func_name}()" elif isinstance(node, ast.Attribute): if node.attr in self.FORBIDDEN_DUNDERS: return f"Forbidden attribute access: .{node.attr}" return None def _check_module(self, module_name: str) -> str | None: """Проверить, разрешён ли модуль. Чистая функция.""" if module_name in self.config.forbidden_modules: return f"Forbidden module: {module_name}" if self.config.allowed_modules: if module_name not in self.config.allowed_modules: return ( f"Module not in whitelist: {module_name}. " f"Allowed: {', '.join(sorted(self.config.allowed_modules))}" ) return None # ============================================================================ # REPL Session # ============================================================================ class ReplSession: """ Одна REPL сессия с persistent subprocess. Управляет: - Долгоживущим worker процессом - Историей команд - Idle timeout - Рабочей директорией """ def __init__( self, session_id: str, config: PythonReplConfig, validator: ReplCodeValidator, ): self.session_id = session_id self.config = config self.validator = validator self.created_at = datetime.now() self.last_active_at = datetime.now() self.history: list[HistoryEntry] = [] self._command_counter = 0 # Subprocess state self._process: asyncio.subprocess.Process | None = None self._work_dir: Path | None = None self._worker_script_path: Path | None = None self._ready = False # Request/response tracking self._request_id = 0 self._pending: dict[int, asyncio.Future] = {} # Background reader task self._reader_task: asyncio.Task | None = None async def start(self) -> None: """Запустить worker subprocess.""" if self._process is not None: return # Создаём work directory if self.config.working_dir: base = Path(self.config.working_dir) base.mkdir(parents=True, exist_ok=True) self._work_dir = base / f"repl_{self.session_id}" else: self._work_dir = Path( tempfile.mkdtemp(prefix=f"repl_{self.session_id}_") ) self._work_dir.mkdir(parents=True, exist_ok=True) # Записываем worker script self._worker_script_path = self._work_dir / "worker.py" max_memory_bytes = self.config.max_memory_mb * 1024 * 1024 worker_code = _generate_worker_script( allowed_modules=self.config.allowed_modules, forbidden_modules=self.config.forbidden_modules, forbidden_builtins=self.config.forbidden_builtins, max_memory_bytes=int(max_memory_bytes), ) self._worker_script_path.write_text(worker_code, encoding="utf-8") # Запускаем subprocess env = os.environ.copy() env["PYTHONDONTWRITEBYTECODE"] = "1" env["PYTHONNOUSERSITE"] = "1" env.pop("PYTHONPATH", None) env.pop("PYTHONSTARTUP", None) try: self._process = await asyncio.create_subprocess_exec( self.config.python_executable, str(self._worker_script_path), stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, cwd=str(self._work_dir), env=env, ) except FileNotFoundError as e: raise ReplWorkerError( f"Python executable not found: {self.config.python_executable}" ) from e except Exception as e: raise ReplWorkerError(f"Failed to start worker: {e}") from e # Запускаем background reader self._reader_task = asyncio.create_task(self._reader_loop()) # Ждём ready signal try: await asyncio.wait_for(self._wait_ready(), timeout=10.0) except asyncio.TimeoutError: await self.shutdown() raise ReplWorkerError("Worker did not send ready signal in time") async def _wait_ready(self) -> None: """Подождать пока worker пришлёт ready signal.""" while not self._ready: await asyncio.sleep(0.05) if self._process is None or self._process.returncode is not None: raise ReplProcessCrashedError("Worker process died before ready") async def _reader_loop(self) -> None: """Background loop — читает JSON ответы от worker.""" if self._process is None or self._process.stdout is None: return try: while True: line_bytes = await self._process.stdout.readline() if not line_bytes: # EOF break line = line_bytes.decode("utf-8", errors="replace").strip() if not line: continue try: message = json.loads(line) except json.JSONDecodeError: logger.warning(f"Worker sent invalid JSON: {line[:200]}") continue # Ready signal if message.get("ready"): self._ready = True continue # Response — route to pending future msg_id = message.get("id") if msg_id is not None and msg_id in self._pending: future = self._pending.pop(msg_id) if not future.done(): future.set_result(message) else: logger.debug(f"Worker message with no pending request: {message}") except asyncio.CancelledError: return except Exception as e: logger.exception(f"Reader loop crashed: {e}") async def _send_request(self, request: dict) -> dict: """Отправить JSON запрос worker и дождаться ответа.""" if self._process is None or self._process.stdin is None: raise ReplProcessCrashedError("Worker process is not running") if self._process.returncode is not None: raise ReplProcessCrashedError( f"Worker process exited with code {self._process.returncode}" ) # Уникальный ID self._request_id += 1 request_id = self._request_id request["id"] = request_id # Создаём future для ответа loop = asyncio.get_running_loop() future: asyncio.Future = loop.create_future() self._pending[request_id] = future try: # Пишем запрос line = json.dumps(request, default=str) + "\n" self._process.stdin.write(line.encode("utf-8")) await self._process.stdin.drain() # Ждём ответ с timeout try: response = await asyncio.wait_for( future, timeout=self.config.command_timeout_seconds ) return response except asyncio.TimeoutError: # Timeout — убираем из pending self._pending.pop(request_id, None) raise ReplTimeoutError( f"Command timed out after {self.config.command_timeout_seconds}s" ) except (ConnectionResetError, BrokenPipeError) as e: self._pending.pop(request_id, None) raise ReplProcessCrashedError(f"Lost connection to worker: {e}") from e async def execute(self, code: str) -> ExecutionOutput: """ Выполнить Python код в сессии. Args: code: Python код для выполнения. Returns: ExecutionOutput с stdout/stderr/результатом. """ if not self._ready: await self.start() # Pre-validation violations = self.validator.validate(code) if violations: raise ReplValidationError( f"Code validation failed: {len(violations)} violation(s)", violations=violations, ) # Проверка длины if len(code) > self.config.max_code_length: raise ReplError( f"Code too long: {len(code)} > {self.config.max_code_length}" ) self.last_active_at = datetime.now() self._command_counter += 1 # Отправляем запрос response = await self._send_request({"command": "execute", "code": code}) # Обрезаем вывод stdout = response.get("stdout", "") or "" stderr = response.get("stderr", "") or "" if len(stdout) > self.config.max_output_length: stdout = ( stdout[: self.config.max_output_length] + f"\n... [truncated, {len(stdout)} chars total]" ) if len(stderr) > self.config.max_output_length: stderr = ( stderr[: self.config.max_output_length] + f"\n... [truncated, {len(stderr)} chars total]" ) output = ExecutionOutput( success=bool(response.get("success", False)), stdout=stdout, stderr=stderr, result_repr=response.get("result_repr"), error_type=response.get("error_type"), error_message=response.get("error_message"), traceback=response.get("traceback"), execution_time_ms=float(response.get("execution_time_ms", 0.0)), ) # Добавляем в историю summary = self._make_history_summary(output) self.history.append( HistoryEntry( index=self._command_counter, code=code, timestamp=datetime.now(), success=output.success, summary=summary, ) ) # Ограничиваем размер истории if len(self.history) > self.config.max_history_size: self.history = self.history[-self.config.max_history_size :] return output def _make_history_summary(self, output: ExecutionOutput) -> str: """Создать краткое описание результата для истории. Чистая функция.""" if not output.success: return f"ERROR: {output.error_type}: {output.error_message}" parts = [] if output.stdout: first_line = output.stdout.strip().split("\n")[0][:80] parts.append(f"stdout: {first_line}") if output.result_repr and output.result_repr != "None": repr_short = output.result_repr[:80] parts.append(f"result: {repr_short}") if not parts: return "(no output)" return "; ".join(parts) async def list_variables(self) -> list[str]: """Получить список определённых переменных.""" if not self._ready: await self.start() response = await self._send_request({"command": "list_variables"}) if not response.get("success"): raise ReplWorkerError( f"list_variables failed: {response.get('error_message')}" ) variables = response.get("variables", []) return [str(v) for v in variables if isinstance(v, str)] async def inspect_variable(self, name: str) -> dict[str, Any]: """Получить информацию о переменной.""" if not self._ready: await self.start() response = await self._send_request( {"command": "inspect", "name": name} ) return response async def reset(self) -> None: """Сбросить состояние сессии (очистить переменные).""" if not self._ready: return response = await self._send_request({"command": "reset"}) if not response.get("success"): raise ReplWorkerError(f"Reset failed: {response.get('error_message')}") self.history.clear() self._command_counter = 0 async def shutdown(self) -> None: """Завершить worker subprocess.""" # Отменяем reader task if self._reader_task is not None: self._reader_task.cancel() try: await self._reader_task except (asyncio.CancelledError, Exception): pass self._reader_task = None # Отправляем shutdown если процесс жив if ( self._process is not None and self._process.returncode is None and self._process.stdin is not None ): try: shutdown_request = json.dumps( {"id": -1, "command": "shutdown"} ) + "\n" self._process.stdin.write(shutdown_request.encode("utf-8")) await self._process.stdin.drain() # Даём время на graceful shutdown try: await asyncio.wait_for(self._process.wait(), timeout=2.0) except asyncio.TimeoutError: self._process.terminate() try: await asyncio.wait_for(self._process.wait(), timeout=2.0) except asyncio.TimeoutError: self._process.kill() await self._process.wait() except Exception as e: logger.warning(f"Error during shutdown: {e}") if self._process.returncode is None: try: self._process.kill() await self._process.wait() except Exception: pass # Resolve все pending futures с ошибкой for future in self._pending.values(): if not future.done(): future.set_exception( ReplProcessCrashedError("Session shutdown") ) self._pending.clear() self._process = None self._ready = False # Cleanup work directory if self.config.auto_cleanup and self._work_dir is not None: try: shutil.rmtree(self._work_dir, ignore_errors=True) except Exception as e: logger.warning(f"Failed to cleanup {self._work_dir}: {e}") self._work_dir = None def is_alive(self) -> bool: """Проверить, жив ли worker процесс.""" return ( self._process is not None and self._process.returncode is None and self._ready ) def get_idle_seconds(self) -> float: """Получить время простоя в секундах.""" return (datetime.now() - self.last_active_at).total_seconds() async def get_info(self) -> ReplSessionInfo: """Получить информацию о сессии.""" variables_count = 0 if self.is_alive(): try: variables = await self.list_variables() variables_count = len(variables) except Exception: variables_count = -1 # unknown return ReplSessionInfo( session_id=self.session_id, created_at=self.created_at, last_active_at=self.last_active_at, commands_count=len(self.history), variables_count=variables_count, idle_seconds=self.get_idle_seconds(), ) # ============================================================================ # Session Manager # ============================================================================ class SessionManager: """ Менеджер REPL сессий. Управляет: - Созданием/уничтожением сессий - Лимитом количества сессий - Auto-cleanup по idle timeout """ def __init__(self, config: PythonReplConfig): self.config = config self.validator = ReplCodeValidator(config) self._sessions: dict[str, ReplSession] = {} self._cleanup_task: asyncio.Task | None = None async def start_cleanup_loop(self) -> None: """Запустить background cleanup loop.""" if self._cleanup_task is None: self._cleanup_task = asyncio.create_task(self._cleanup_loop()) async def stop_cleanup_loop(self) -> None: """Остановить cleanup loop.""" if self._cleanup_task is not None: self._cleanup_task.cancel() try: await self._cleanup_task except asyncio.CancelledError: pass self._cleanup_task = None async def _cleanup_loop(self) -> None: """Periodically удалять idle сессии.""" try: while True: await asyncio.sleep(30.0) # check every 30s await self._cleanup_idle_sessions() except asyncio.CancelledError: return async def _cleanup_idle_sessions(self) -> None: """Удалить сессии с idle timeout.""" to_remove: list[str] = [] for session_id, session in self._sessions.items(): if session.get_idle_seconds() > self.config.session_idle_timeout_seconds: to_remove.append(session_id) for session_id in to_remove: logger.info(f"Auto-destroying idle session: {session_id}") await self.destroy_session(session_id) async def get_or_create_session( self, session_id: str | None ) -> ReplSession: """ Получить существующую сессию или создать новую. Args: session_id: ID сессии (None = создать новую) Returns: ReplSession instance Raises: ReplSessionNotFoundError: если session_id задан, но не найден ReplError: если достигнут лимит сессий """ await self.start_cleanup_loop() if session_id is not None: session = self._sessions.get(session_id) if session is None: raise ReplSessionNotFoundError( f"Session not found: {session_id}. " f"Active sessions: {list(self._sessions.keys())}" ) # Проверяем что жива if not session.is_alive(): # Перезапускаем await session.start() return session # Создаём новую сессию if len(self._sessions) >= self.config.max_sessions: # Пытаемся освободить место, удалив самую старую idle сессию await self._cleanup_idle_sessions() if len(self._sessions) >= self.config.max_sessions: # Удаляем самую старую oldest_id = min( self._sessions.keys(), key=lambda sid: self._sessions[sid].last_active_at, ) logger.warning( f"Max sessions reached, destroying oldest: {oldest_id}" ) await self.destroy_session(oldest_id) new_id = uuid.uuid4().hex[:12] session = ReplSession(new_id, self.config, self.validator) self._sessions[new_id] = session await session.start() return session async def destroy_session(self, session_id: str) -> bool: """Уничтожить сессию.""" session = self._sessions.pop(session_id, None) if session is None: return False await session.shutdown() return True async def list_sessions(self) -> list[ReplSessionInfo]: """Получить информацию о всех сессиях.""" infos: list[ReplSessionInfo] = [] for session in self._sessions.values(): try: info = await session.get_info() infos.append(info) except Exception as e: logger.warning(f"Failed to get session info: {e}") return infos async def shutdown_all(self) -> None: """Завершить все сессии.""" await self.stop_cleanup_loop() for session in list(self._sessions.values()): await session.shutdown() self._sessions.clear() # ============================================================================ # Python REPL Tool (MCP) # ============================================================================ class PythonReplTool(Tool): """ MCP Tool для stateful выполнения Python кода (REPL). В отличие от code_execution (stateless), этот tool поддерживает persistent сессии — состояние сохраняется между вызовами. Operations: - `execute` — выполнить код в сессии - `reset` — сбросить состояние сессии - `destroy` — уничтожить сессию - `list_sessions` — список активных сессий - `history` — история команд сессии - `inspect` — информация о переменной - `list_variables` — список переменных Workflow: 1. Первый execute без session_id создаёт новую сессию и возвращает session_id 2. Последующие execute с session_id используют ту же сессию 3. Переменные, импорты, функции сохраняются между вызовами 4. Сессии автоматически удаляются после idle timeout (10 минут) """ name: ClassVar[str] = "python_repl" description: ClassVar[str] = ( "Stateful Python REPL — выполнение кода с сохранением состояния между вызовами. " "Поддерживает persistent сессии: переменные, импорты, определения функций " "остаются доступными между вызовами. " "Operations: execute, reset, destroy, list_sessions, history, inspect, list_variables. " "Первый execute без session_id создаёт новую сессию." ) input_schema: ClassVar[dict[str, Any] | None] = { "type": "object", "properties": { "operation": { "type": "string", "enum": [ "execute", "reset", "destroy", "list_sessions", "history", "inspect", "list_variables", ], "description": "Тип операции", }, "code": { "type": "string", "description": "Python код для выполнения (для operation=execute)", }, "session_id": { "type": "string", "description": ( "ID сессии. Если не задан для execute — создаётся новая сессия. " "Для других операций (кроме list_sessions) — обязателен." ), }, "variable_name": { "type": "string", "description": "Имя переменной для inspect", }, }, "required": ["operation"], "additionalProperties": False, } parameters_schema: ClassVar[dict[str, Any] | None] = input_schema examples: ClassVar[list[dict[str, Any]]] = [ { "operation": "execute", "code": "x = 42\nprint(x)", "note": "Создаёт новую сессию, возвращает session_id", }, { "operation": "execute", "session_id": "abc123", "code": "print(x * 2)", "note": "Использует x из предыдущего вызова", }, { "operation": "list_variables", "session_id": "abc123", }, { "operation": "inspect", "session_id": "abc123", "variable_name": "x", }, ] tags: ClassVar[list[str]] = ["python", "repl", "code", "stateful"] is_read_only: ClassVar[bool] = False requires_confirmation: ClassVar[bool] = False def __init__(self, config: PythonReplConfig | None = None): self.config = config or PythonReplConfig() self._manager: SessionManager | None = None def _get_manager(self) -> SessionManager: """Получить или создать SessionManager.""" if self._manager is None: self._manager = SessionManager(self.config) return self._manager async def execute(self, **kwargs: Any) -> ToolResult: """ Выполнить REPL операцию. Args: operation: Тип операции code: Python код (для execute) session_id: ID сессии (опционально для execute, обязательно для других) variable_name: Имя переменной (для inspect) Returns: ToolResult с результатом операции. """ operation_raw = kwargs.get("operation", "") if not isinstance(operation_raw, str) or not operation_raw: return ToolResult.failure("Parameter 'operation' is required") operation = operation_raw.strip().lower() valid_ops = set(self.input_schema["properties"]["operation"]["enum"]) # type: ignore[index] if operation not in valid_ops: return ToolResult.failure( f"Unknown operation: {operation}. " f"Valid: {', '.join(sorted(valid_ops))}" ) session_id = kwargs.get("session_id") if session_id is not None and not isinstance(session_id, str): return ToolResult.failure("session_id must be a string") try: manager = self._get_manager() if operation == "execute": return await self._handle_execute(manager, kwargs) elif operation == "reset": return await self._handle_reset(manager, session_id) elif operation == "destroy": return await self._handle_destroy(manager, session_id) elif operation == "list_sessions": return await self._handle_list_sessions(manager) elif operation == "history": return await self._handle_history(manager, session_id) elif operation == "inspect": return await self._handle_inspect(manager, session_id, kwargs) elif operation == "list_variables": return await self._handle_list_variables(manager, session_id) else: return ToolResult.failure(f"Unhandled operation: {operation}") except ReplSessionNotFoundError as e: return ToolResult.failure( str(e), metadata={"operation": operation, "error_type": "session_not_found"}, ) except ReplValidationError as e: violations_text = "\n".join(f" - {v}" for v in e.violations) return ToolResult.failure( f"Code validation failed:\n{violations_text}", metadata={ "operation": operation, "error_type": "validation", "violations": e.violations, }, ) except ReplTimeoutError as e: return ToolResult.failure( str(e), metadata={"operation": operation, "error_type": "timeout"}, ) except ReplProcessCrashedError as e: return ToolResult.failure( f"Worker process crashed: {e}", metadata={"operation": operation, "error_type": "process_crashed"}, ) except ReplWorkerError as e: return ToolResult.failure( f"Worker error: {e}", metadata={"operation": operation, "error_type": "worker"}, ) except ReplError as e: return ToolResult.failure( f"REPL error: {e}", metadata={"operation": operation}, ) except Exception as e: logger.exception(f"Unexpected error in python_repl tool: {e}") return ToolResult.failure( f"Unexpected error: {type(e).__name__}: {e}", metadata={"operation": operation}, ) async def _handle_execute( self, manager: SessionManager, kwargs: dict[str, Any] ) -> ToolResult: """Обработать operation=execute.""" code = kwargs.get("code") if not isinstance(code, str) or not code.strip(): return ToolResult.failure( "Parameter 'code' (non-empty string) is required for operation 'execute'" ) session_id = kwargs.get("session_id") session = await manager.get_or_create_session(session_id) output = await session.execute(code) # Формируем текстовый вывод content_items: list[Content] = [] lines = [f"# Python REPL — Execute\n"] lines.append(f"**Session ID:** `{session.session_id}`") lines.append(f"**Execution time:** {output.execution_time_ms:.1f}ms\n") lines.append("## Code\n```python\n" + code + "\n```") if output.stdout: lines.append("\n## Stdout\n```\n" + output.stdout + "\n```") if output.result_repr is not None and output.result_repr != "None": lines.append(f"\n## Result\n```python\n{output.result_repr}\n```") if not output.success: lines.append( f"\n## Error: {output.error_type}\n" f"```\n{output.error_message}\n```" ) if output.traceback: lines.append("\n### Traceback\n```python\n" + output.traceback + "\n```") if output.stderr: lines.append("\n## Stderr\n```\n" + output.stderr + "\n```") content_items.append(TextContent(text="\n".join(lines))) metadata = { "operation": "execute", "session_id": session.session_id, "created_new_session": session_id is None, **output.to_dict(), } if output.success: return ToolResult.success_result(content_items, metadata=metadata) else: return ToolResult.failure( error=f"{output.error_type}: {output.error_message}", content=content_items, metadata=metadata, ) async def _handle_reset( self, manager: SessionManager, session_id: str | None ) -> ToolResult: """Обработать operation=reset.""" if not session_id: return ToolResult.failure( "Parameter 'session_id' is required for operation 'reset'" ) session = await manager.get_or_create_session(session_id) await session.reset() text = ( f"# Session Reset\n\n" f"**Session ID:** `{session_id}`\n\n" f"All variables cleared. History reset." ) return ToolResult.success_result( [TextContent(text=text)], metadata={ "operation": "reset", "session_id": session_id, }, ) async def _handle_destroy( self, manager: SessionManager, session_id: str | None ) -> ToolResult: """Обработать operation=destroy.""" if not session_id: return ToolResult.failure( "Parameter 'session_id' is required for operation 'destroy'" ) destroyed = await manager.destroy_session(session_id) if destroyed: text = ( f"# Session Destroyed\n\n" f"**Session ID:** `{session_id}`\n\n" f"Worker process terminated, resources released." ) return ToolResult.success_result( [TextContent(text=text)], metadata={"operation": "destroy", "session_id": session_id, "destroyed": True}, ) else: return ToolResult.failure( f"Session not found: {session_id}", metadata={"operation": "destroy", "session_id": session_id}, ) async def _handle_list_sessions(self, manager: SessionManager) -> ToolResult: """Обработать operation=list_sessions.""" infos = await manager.list_sessions() lines = [ "# Active REPL Sessions\n", f"**Total:** {len(infos)}\n", ] if not infos: lines.append("*(no active sessions)*") else: for info in infos: lines.append(f"## Session `{info.session_id}`") lines.append( f"- **Created:** {info.created_at.isoformat()}" ) lines.append( f"- **Last active:** {info.last_active_at.isoformat()}" ) lines.append(f"- **Idle:** {info.idle_seconds:.0f}s") lines.append(f"- **Commands:** {info.commands_count}") lines.append(f"- **Variables:** {info.variables_count}") lines.append("") return ToolResult.success_result( [TextContent(text="\n".join(lines))], metadata={ "operation": "list_sessions", "sessions": [info.to_dict() for info in infos], "count": len(infos), }, ) async def _handle_history( self, manager: SessionManager, session_id: str | None ) -> ToolResult: """Обработать operation=history.""" if not session_id: return ToolResult.failure( "Parameter 'session_id' is required for operation 'history'" ) session = await manager.get_or_create_session(session_id) lines = [ f"# Session History\n", f"**Session ID:** `{session_id}`\n", f"**Commands:** {len(session.history)}\n", ] if not session.history: lines.append("*(no commands executed yet)*") else: for entry in session.history: status = "✅" if entry.success else "❌" lines.append( f"### [{entry.index}] {status} {entry.timestamp.strftime('%H:%M:%S')}" ) # Обрезаем длинный код code_preview = entry.code if len(code_preview) > 500: code_preview = code_preview[:500] + "... <truncated>" lines.append(f"```python\n{code_preview}\n```") lines.append(f"**→** {entry.summary}") lines.append("") return ToolResult.success_result( [TextContent(text="\n".join(lines))], metadata={ "operation": "history", "session_id": session_id, "history": [e.to_dict() for e in session.history], "count": len(session.history), }, ) async def _handle_inspect( self, manager: SessionManager, session_id: str | None, kwargs: dict[str, Any], ) -> ToolResult: """Обработать operation=inspect.""" if not session_id: return ToolResult.failure( "Parameter 'session_id' is required for operation 'inspect'" ) variable_name = kwargs.get("variable_name") if not isinstance(variable_name, str) or not variable_name: return ToolResult.failure( "Parameter 'variable_name' is required for operation 'inspect'" ) session = await manager.get_or_create_session(session_id) info = await session.inspect_variable(variable_name) if not info.get("found"): return ToolResult.failure( f"Variable '{variable_name}' not found in session {session_id}", metadata={ "operation": "inspect", "session_id": session_id, "variable_name": variable_name, }, ) text = ( f"# Variable: `{variable_name}`\n\n" f"**Session:** `{session_id}`\n\n" f"**Type:** `{info.get('type', 'unknown')}`\n\n" f"**Value:**\n```python\n{info.get('repr', '')}\n```" ) return ToolResult.success_result( [TextContent(text=text)], metadata={ "operation": "inspect", "session_id": session_id, **info, }, ) async def _handle_list_variables( self, manager: SessionManager, session_id: str | None ) -> ToolResult: """Обработать operation=list_variables.""" if not session_id: return ToolResult.failure( "Parameter 'session_id' is required for operation 'list_variables'" ) session = await manager.get_or_create_session(session_id) variables = await session.list_variables() lines = [ f"# Session Variables\n", f"**Session ID:** `{session_id}`\n", f"**Count:** {len(variables)}\n", ] if not variables: lines.append("*(no variables defined)*") else: for var in sorted(variables): lines.append(f"- `{var}`") return ToolResult.success_result( [TextContent(text="\n".join(lines))], metadata={ "operation": "list_variables", "session_id": session_id, "variables": variables, "count": len(variables), }, ) async def shutdown(self) -> None: """Завершить все сессии (вызывать при shutdown приложения).""" if self._manager is not None: await self._manager.shutdown_all() # ============================================================================ # Self-test # ============================================================================ if __name__ == "__main__": import asyncio async def _self_test() -> None: """Проверка корректности python_repl tool.""" # === Helper functions для type-safe доступа к metadata === def _meta_bool(metadata: dict[str, Any] | None, key: str) -> bool: """Безопасно извлечь bool из metadata. Чистая функция.""" if not isinstance(metadata, dict): return False value = metadata.get(key) return value is True def _meta_int(metadata: dict[str, Any] | None, key: str, default: int = 0) -> int: """Безопасно извлечь int из metadata. Чистая функция.""" if not isinstance(metadata, dict): return default value = metadata.get(key) return int(value) if isinstance(value, (int, float)) else default def _meta_list_str(metadata: dict[str, Any] | None, key: str) -> list[str]: """Безопасно извлечь list[str] из metadata. Чистая функция.""" if not isinstance(metadata, dict): return [] value = metadata.get(key) if isinstance(value, list): return [str(x) for x in value if isinstance(x, str)] return [] def _meta_str(metadata: dict[str, Any] | None, key: str) -> str | None: """Безопасно извлечь str из metadata. Чистая функция.""" if not isinstance(metadata, dict): return None value = metadata.get(key) return str(value) if isinstance(value, str) else None config = PythonReplConfig( session_idle_timeout_seconds=60.0, max_sessions=5, ) tool = PythonReplTool(config) print("🐍 Python REPL Tool self-test:\n") passed = 0 failed = 0 def check(name: str, ok: bool) -> None: nonlocal passed, failed status = "✅" if ok else "❌" print(f" {status} {name}") if ok: passed += 1 else: failed += 1 # === Создание сессии и первый execute === print("=== Session creation ===") result = await tool.execute(operation="execute", code="x = 42\nprint(x)") ok = ( result.is_success() and _meta_bool(result.metadata, "created_new_session") and "42" in result.get_text() ) check("Create session + execute", ok) session_id = _meta_str(result.metadata, "session_id") assert session_id is not None, "session_id should be set" # === Использование состояния из предыдущего вызова === print("\n=== State persistence ===") result = await tool.execute( operation="execute", session_id=session_id, code="print(x * 2)", ) ok = result.is_success() and "84" in result.get_text() check("State preserved between calls", ok) # === Определение функции === result = await tool.execute( operation="execute", session_id=session_id, code="def square(n): return n ** 2", ) ok = result.is_success() check("Define function", ok) result = await tool.execute( operation="execute", session_id=session_id, code="print(square(7))", ) ok = result.is_success() and "49" in result.get_text() check("Use defined function", ok) # === Импорт модуля === print("\n=== Imports ===") result = await tool.execute( operation="execute", session_id=session_id, code="import math\nprint(math.pi)", ) ok = result.is_success() and "3.14" in result.get_text() check("Import safe module (math)", ok) result = await tool.execute( operation="execute", session_id=session_id, code="import os", ) ok = result.is_failure() and "forbidden" in result.get_text().lower() check("Block forbidden module (os)", ok) # === list_variables === print("\n=== Variables inspection ===") result = await tool.execute( operation="list_variables", session_id=session_id, ) variables = _meta_list_str(result.metadata, "variables") ok = ( result.is_success() and "x" in variables and "square" in variables ) check("list_variables", ok) # === inspect === result = await tool.execute( operation="inspect", session_id=session_id, variable_name="x", ) ok = ( result.is_success() and _meta_str(result.metadata, "type") == "int" and _meta_str(result.metadata, "repr") == "42" ) check("inspect variable", ok) # === history === print("\n=== History ===") result = await tool.execute( operation="history", session_id=session_id, ) ok = ( result.is_success() and _meta_int(result.metadata, "count", 0) >= 4 ) check("history", ok) # === list_sessions === print("\n=== Session management ===") result = await tool.execute(operation="list_sessions") ok = ( result.is_success() and _meta_int(result.metadata, "count", 0) >= 1 ) check("list_sessions", ok) # === reset === result = await tool.execute(operation="reset", session_id=session_id) ok = result.is_success() check("reset", ok) result = await tool.execute( operation="execute", session_id=session_id, code="print(x)", # x должна быть undefined после reset ) ok = result.is_failure() and "NameError" in result.get_text() check("reset clears state", ok) # === Error handling === print("\n=== Error handling ===") result = await tool.execute( operation="execute", session_id=session_id, code="1 / 0", ) ok = result.is_failure() and "ZeroDivisionError" in result.get_text() check("Runtime error handling", ok) result = await tool.execute( operation="execute", session_id=session_id, code="this is not valid python syntax", ) ok = result.is_failure() and "Syntax" in result.get_text() check("Syntax error handling", ok) # === Security === print("\n=== Security ===") result = await tool.execute( operation="execute", session_id=session_id, code="__import__('os').system('echo pwned')", ) ok = result.is_failure() check("Block __import__", ok) result = await tool.execute( operation="execute", session_id=session_id, code="open('/etc/passwd').read()", ) ok = result.is_failure() check("Block open()", ok) result = await tool.execute( operation="execute", session_id=session_id, code="eval('1+1')", ) ok = result.is_failure() check("Block eval()", ok) result = await tool.execute( operation="execute", session_id=session_id, code="().__class__.__bases__[0].__subclasses__()", ) ok = result.is_failure() check("Block dunder escape", ok) # === Session not found === print("\n=== Edge cases ===") result = await tool.execute( operation="execute", session_id="nonexistent_session_id", code="print(1)", ) ok = result.is_failure() and "not found" in result.get_text().lower() check("Session not found error", ok) # === destroy === result = await tool.execute( operation="destroy", session_id=session_id, ) ok = ( result.is_success() and _meta_bool(result.metadata, "destroyed") ) check("destroy session", ok) result = await tool.execute( operation="execute", session_id=session_id, code="print(1)", ) ok = result.is_failure() and "not found" in result.get_text().lower() check("Destroyed session not usable", ok) # === MCP Schema === print("\n📝 MCP Tool Schema:") print(json.dumps(tool.to_mcp_tool(), indent=2, ensure_ascii=False)) print(f"\n📊 Results: {passed} passed, {failed} failed") # Cleanup await tool.shutdown() asyncio.run(_self_test())