Add python_exec tool: run Python code in subprocess with ctx injection and timeout

Executes code in a fresh python3 subprocess (max 30s, default 10s).
Pre-injects a 'ctx' dict with: chat_id, message_id, message_text, args,
from_user, replied_message. stdout+stderr returned; output capped at 3000 chars.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-05-18 14:16:12 +02:00
parent 95cac39fcf
commit c85af7c54d
2 changed files with 115 additions and 3 deletions

View File

@@ -1,10 +1,12 @@
from __future__ import annotations from __future__ import annotations
import asyncio
from dataclasses import dataclass from dataclasses import dataclass
import inspect import inspect
import json import json
import logging import logging
from pathlib import Path from pathlib import Path
import sys
import tempfile import tempfile
from typing import TYPE_CHECKING, Any, Awaitable, Callable from typing import TYPE_CHECKING, Any, Awaitable, Callable
@@ -19,7 +21,7 @@ from app.ai.image_understand import (
) )
from app.ai.parser import AIInvocation from app.ai.parser import AIInvocation
from app.ai.prompts import PromptBook from app.ai.prompts import PromptBook
from app.ai.tools import FILE_TOOL_DISPATCH, FILE_TOOL_SCHEMAS, mcp_tools_to_openai_schemas from app.ai.tools import FILE_TOOL_DISPATCH, FILE_TOOL_SCHEMAS, PYTHON_EXEC_SCHEMA, mcp_tools_to_openai_schemas
from app.config import AIConfig, AIModelConfig, AIProviderConfig from app.config import AIConfig, AIModelConfig, AIProviderConfig
if TYPE_CHECKING: if TYPE_CHECKING:
@@ -327,11 +329,92 @@ class AIService:
async def _build_tool_list(self) -> list[dict]: async def _build_tool_list(self) -> list[dict]:
mcp_schemas = await self._fetch_mcp_tool_schemas() mcp_schemas = await self._fetch_mcp_tool_schemas()
return FILE_TOOL_SCHEMAS + mcp_schemas return FILE_TOOL_SCHEMAS + [PYTHON_EXEC_SCHEMA] + mcp_schemas
def _get_allowed_roots(self) -> tuple[Path, ...]: def _get_allowed_roots(self) -> tuple[Path, ...]:
return (self.config.workspace_dir, self.config.prompt.root) return (self.config.workspace_dir, self.config.prompt.root)
def _build_exec_context(self, ctx: "CommandContext") -> dict[str, Any]:
actor = ctx.actor
replied = ctx.replied_message
replied_user = ctx.replied_user
return {
"chat_id": getattr(getattr(ctx.message, "chat", None), "id", None),
"message_id": getattr(ctx.message, "id", None),
"message_text": getattr(ctx.message, "text", None) or "",
"args": ctx.args,
"from_user": {
"id": getattr(actor, "id", None),
"first_name": getattr(actor, "first_name", None),
"last_name": getattr(actor, "last_name", None),
"username": getattr(actor, "username", None),
} if actor else None,
"replied_message": {
"id": getattr(replied, "id", None),
"text": getattr(replied, "text", None) or getattr(replied, "caption", None) or "",
"from_user": {
"id": getattr(replied_user, "id", None),
"first_name": getattr(replied_user, "first_name", None),
"username": getattr(replied_user, "username", None),
} if replied_user else None,
} if replied else None,
}
async def _execute_python_exec(
self,
arguments: dict,
ctx: "CommandContext",
) -> str:
code = str(arguments.get("code", "")).strip()
if not code:
raise ValueError("python_exec: 'code' is required")
raw_timeout = arguments.get("timeout", 10)
try:
timeout = float(raw_timeout)
except (TypeError, ValueError):
timeout = 10.0
timeout = min(max(timeout, 1.0), 30.0)
context = self._build_exec_context(ctx)
context_json = json.dumps(context, ensure_ascii=False, default=str)
context_line = f"import json as _json; ctx = _json.loads({context_json!r})\n"
full_code = (context_line + code).encode("utf-8")
try:
proc = await asyncio.create_subprocess_exec(
sys.executable, "-",
stdin=asyncio.subprocess.PIPE,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
)
try:
stdout, stderr = await asyncio.wait_for(
proc.communicate(input=full_code), timeout=timeout
)
except asyncio.TimeoutError:
proc.kill()
await proc.wait()
return f"Execution timed out after {timeout:.0f}s"
except OSError as exc:
raise ValueError(f"python_exec: failed to start interpreter: {exc}") from exc
out = stdout.decode("utf-8", errors="replace").rstrip()
err = stderr.decode("utf-8", errors="replace").rstrip()
parts: list[str] = []
if out:
parts.append(out)
if err:
parts.append(f"[stderr]\n{err}")
if not parts:
parts.append(f"[exit {proc.returncode}]" if proc.returncode else "(no output)")
result = "\n".join(parts)
if len(result) > 3000:
result = result[:3000] + "\n…(truncated)"
return result
async def _execute_send_file( async def _execute_send_file(
self, self,
arguments: dict, arguments: dict,
@@ -373,7 +456,9 @@ class AIService:
return f"Tool {name!r}: invalid JSON arguments: {exc}" return f"Tool {name!r}: invalid JSON arguments: {exc}"
try: try:
if name == "send_file": if name == "python_exec":
result = await self._execute_python_exec(arguments, ctx)
elif name == "send_file":
result = await self._execute_send_file(arguments, ctx, allowed_roots) result = await self._execute_send_file(arguments, ctx, allowed_roots)
elif name in FILE_TOOL_DISPATCH: elif name in FILE_TOOL_DISPATCH:
result = await FILE_TOOL_DISPATCH[name](arguments, allowed_roots) result = await FILE_TOOL_DISPATCH[name](arguments, allowed_roots)

View File

@@ -184,6 +184,33 @@ def _walk(
pass pass
PYTHON_EXEC_SCHEMA: dict[str, Any] = {
"type": "function",
"function": {
"name": "python_exec",
"description": (
"Execute Python code and return stdout/stderr output. "
"A 'ctx' dict is pre-injected with message context: "
"chat_id, message_id, message_text, args, from_user, replied_message. "
"Use print() to produce output."
),
"parameters": {
"type": "object",
"properties": {
"code": {
"type": "string",
"description": "Python code to execute",
},
"timeout": {
"type": "number",
"description": "Execution timeout in seconds (default 10, max 30)",
},
},
"required": ["code"],
},
},
}
ToolExecutor = Callable[[dict, tuple[Path, ...]], Awaitable[str]] ToolExecutor = Callable[[dict, tuple[Path, ...]], Awaitable[str]]
FILE_TOOL_DISPATCH: dict[str, ToolExecutor] = { FILE_TOOL_DISPATCH: dict[str, ToolExecutor] = {