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
import asyncio
from dataclasses import dataclass
import inspect
import json
import logging
from pathlib import Path
import sys
import tempfile
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.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
if TYPE_CHECKING:
@@ -327,11 +329,92 @@ class AIService:
async def _build_tool_list(self) -> list[dict]:
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, ...]:
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(
self,
arguments: dict,
@@ -373,7 +456,9 @@ class AIService:
return f"Tool {name!r}: invalid JSON arguments: {exc}"
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)
elif name in FILE_TOOL_DISPATCH:
result = await FILE_TOOL_DISPATCH[name](arguments, allowed_roots)

View File

@@ -184,6 +184,33 @@ def _walk(
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]]
FILE_TOOL_DISPATCH: dict[str, ToolExecutor] = {