diff --git a/astrbot/core/tools/computer_tools/shell.py b/astrbot/core/tools/computer_tools/shell.py index af933e83b..bfde5b8fd 100644 --- a/astrbot/core/tools/computer_tools/shell.py +++ b/astrbot/core/tools/computer_tools/shell.py @@ -1,11 +1,15 @@ import json +import os +import uuid from dataclasses import dataclass, field +from pathlib import Path from astrbot.api import FunctionTool from astrbot.core.agent.run_context import ContextWrapper from astrbot.core.agent.tool import ToolExecResult from astrbot.core.astr_agent_context import AstrAgentContext from astrbot.core.computer.computer_client import get_booter +from astrbot.core.utils.astrbot_path import get_astrbot_system_tmp_path from ..registry import builtin_tool from .util import check_admin_permission, is_local_runtime, workspace_root @@ -15,6 +19,32 @@ _COMPUTER_RUNTIME_TOOL_CONFIG = { } +def _quote_redirect_path(path: str, *, local_runtime: bool) -> str: + if local_runtime and os.name == "nt": + escaped_path = path.replace('"', '""') + else: + escaped_path = path.replace("\\", "\\\\").replace('"', '\\"') + return f'"{escaped_path}"' + + +def _build_background_output_path(*, local_runtime: bool) -> str: + file_name = f"astrbot_shell_stdout_{uuid.uuid4().hex[:8]}.log" + if local_runtime: + output_dir = Path(get_astrbot_system_tmp_path()) / "shell" + output_dir.mkdir(parents=True, exist_ok=True) + return str((output_dir / file_name).resolve(strict=False)) + return f"/tmp/{file_name}" + + +def _redirect_background_stdout_command( + command: str, + *, + output_path: str, + local_runtime: bool, +) -> str: + return f"{command} > {_quote_redirect_path(output_path, local_runtime=local_runtime)} 2>&1" + + @builtin_tool(config=_COMPUTER_RUNTIME_TOOL_CONFIG) @dataclass class ExecuteShellTool(FunctionTool): @@ -28,14 +58,18 @@ class ExecuteShellTool(FunctionTool): "type": "string", "description": "The shell command to execute in the current runtime shell (for example, cmd.exe on Windows). Equal to 'cd {working_dir} && {your_command}'.", }, + "timeout": { + "type": "integer", + "description": "Optional timeout in seconds for the command execution.", + }, "background": { "type": "boolean", - "description": "Whether to run the command in the background.", + "description": "Run the command in the background. Use the file read tool to read the output later.", "default": False, }, "env": { "type": "object", - "description": "Optional environment variables to set for the file creation process.", + "description": "Optional environment variables to set.", "additionalProperties": {"type": "string"}, "default": {}, }, @@ -50,6 +84,7 @@ class ExecuteShellTool(FunctionTool): command: str, background: bool = False, env: dict = {}, + timeout: int | None = None, ) -> ToolExecResult: if permission_error := check_admin_permission(context, "Shell execution"): return permission_error @@ -67,12 +102,31 @@ class ExecuteShellTool(FunctionTool): current_workspace_root.mkdir(parents=True, exist_ok=True) cwd = str(current_workspace_root) + stdout_file: str | None = None + if background: + local_runtime = is_local_runtime(context) + stdout_file = _build_background_output_path( + local_runtime=local_runtime, + ) + command = _redirect_background_stdout_command( + command, + output_path=stdout_file, + local_runtime=local_runtime, + ) + result = await sb.shell.exec( command, cwd=cwd, background=background, env=env, + timeout=timeout or context.tool_call_timeout or 30, ) + if stdout_file: + result["stdout_file"] = stdout_file + result["stdout"] = ( + f"Command is running in the background. stdout/stderr is being " + f"written to `{stdout_file}`. Use astrbot_file_read_tool to read it." + ) return json.dumps(result, ensure_ascii=False) except Exception as e: return f"Error executing command: {str(e)}" diff --git a/tests/unit/test_computer_shell_tool.py b/tests/unit/test_computer_shell_tool.py new file mode 100644 index 000000000..d8dfe4f8c --- /dev/null +++ b/tests/unit/test_computer_shell_tool.py @@ -0,0 +1,15 @@ +from astrbot.core.tools.computer_tools.shell import ( + _redirect_background_stdout_command, +) + + +def test_background_shell_command_redirects_output_to_file(): + command = _redirect_background_stdout_command( + "python -c 'print(123)'", + output_path="/tmp/astrbot shell output.log", + local_runtime=False, + ) + + assert command == ( + "python -c 'print(123)' > \"/tmp/astrbot shell output.log\" 2>&1" + )