diff --git a/astrbot/core/tools/computer_tools/python.py b/astrbot/core/tools/computer_tools/python.py index e0bb6c9de6..be909f6d26 100644 --- a/astrbot/core/tools/computer_tools/python.py +++ b/astrbot/core/tools/computer_tools/python.py @@ -33,6 +33,11 @@ "description": "Whether to suppress the output of the code execution.", "default": False, }, + "timeout": { + "type": "integer", + "description": "Optional timeout in seconds for code execution.", + "default": 30, + }, }, "required": ["code"], } @@ -77,7 +82,11 @@ class PythonTool(FunctionTool): parameters: dict = field(default_factory=lambda: param_schema) async def call( - self, context: ContextWrapper[AstrAgentContext], code: str, silent: bool = False + self, + context: ContextWrapper[AstrAgentContext], + code: str, + silent: bool = False, + timeout: int = 30, ) -> ToolExecResult: if permission_error := check_admin_permission(context, "Python execution"): return permission_error @@ -85,8 +94,17 @@ async def call( context.context.context, context.context.event.unified_msg_origin, ) + effective_timeout = ( + min(timeout, context.tool_call_timeout) + if timeout > 0 + else context.tool_call_timeout + ) try: - result = await sb.python.exec(code, silent=silent) + result = await sb.python.exec( + code, + timeout=effective_timeout, + silent=silent, + ) return await handle_result(result, context.context.event) except Exception as e: return f"Error executing code: {str(e)}" @@ -104,13 +122,26 @@ class LocalPythonTool(FunctionTool): parameters: dict = field(default_factory=lambda: param_schema) async def call( - self, context: ContextWrapper[AstrAgentContext], code: str, silent: bool = False + self, + context: ContextWrapper[AstrAgentContext], + code: str, + silent: bool = False, + timeout: int = 30, ) -> ToolExecResult: if permission_error := check_admin_permission(context, "Python execution"): return permission_error sb = get_local_booter() + effective_timeout = ( + min(timeout, context.tool_call_timeout) + if timeout > 0 + else context.tool_call_timeout + ) try: - result = await sb.python.exec(code, silent=silent) + result = await sb.python.exec( + code, + timeout=effective_timeout, + silent=silent, + ) return await handle_result(result, context.context.event) except Exception as e: return f"Error executing code: {str(e)}"