mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-16 17:47:06 +08:00
543 lines
16 KiB
Python
543 lines
16 KiB
Python
import json
|
|
|
|
import pytest
|
|
|
|
from astrbot.core import sp
|
|
from astrbot.core.agent.run_context import ContextWrapper
|
|
from astrbot.core.agent.tool import FunctionTool, ToolExecResult
|
|
from astrbot.core.astr_agent_context import AstrAgentContext
|
|
from astrbot.core.computer.computer_tool_provider import get_all_tools
|
|
from astrbot.core.provider.func_tool_manager import FunctionToolManager
|
|
from astrbot.core.tools.computer_tools.shell import ExecuteShellTool
|
|
from astrbot.core.tools.message_tools import SendMessageToUserTool
|
|
from astrbot.core.tools.web_search_tools import (
|
|
FirecrawlExtractWebPageTool,
|
|
FirecrawlWebSearchTool,
|
|
)
|
|
|
|
|
|
def _make_fake_wrapper_class():
|
|
class FakeConfig:
|
|
def get_config(self, umo):
|
|
return {"provider_settings": {"computer_use_runtime": "sandbox"}}
|
|
|
|
class FakeEvent:
|
|
unified_msg_origin = "umo"
|
|
role = "admin"
|
|
|
|
class FakeAstrContext:
|
|
context = FakeConfig()
|
|
event = FakeEvent()
|
|
|
|
class FakeWrapper(ContextWrapper[AstrAgentContext]):
|
|
def __init__(self):
|
|
self.context = FakeAstrContext() # type: ignore[assignment]
|
|
self.messages = []
|
|
|
|
return FakeWrapper
|
|
|
|
|
|
def test_computer_tools_provider_returns_tools():
|
|
"""ComputerToolProvider should return a list of computer tools."""
|
|
tools = get_all_tools()
|
|
|
|
assert len(tools) > 0
|
|
names = [t.name for t in tools]
|
|
assert "astrbot_execute_shell" in names
|
|
|
|
|
|
def test_register_internal_tools_adds_tools_to_manager():
|
|
"""register_internal_tools should add computer tools to the manager."""
|
|
manager = FunctionToolManager()
|
|
|
|
tool = manager.get_func("astrbot_execute_shell")
|
|
assert tool is not None
|
|
assert tool.name == "astrbot_execute_shell"
|
|
|
|
|
|
def test_manager_func_list_starts_empty():
|
|
"""func_list should start empty in new architecture."""
|
|
manager = FunctionToolManager()
|
|
|
|
assert manager.func_list == []
|
|
|
|
|
|
def test_register_internal_tools_does_not_duplicate():
|
|
"""Calling register_internal_tools twice should not duplicate tools."""
|
|
manager = FunctionToolManager()
|
|
|
|
first_tool = manager.get_func("astrbot_execute_shell")
|
|
assert first_tool is not None
|
|
|
|
# Should still have the same tool (not duplicated)
|
|
second_tool = manager.get_func("astrbot_execute_shell")
|
|
assert second_tool is first_tool
|
|
assert first_tool is not None
|
|
assert first_tool.parameters["properties"]["background"]["default"] is False # type: ignore[union-attr]
|
|
assert manager.is_builtin_tool("astrbot_execute_shell") is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_shell_defaults_to_foreground(monkeypatch):
|
|
from astrbot.core.tools.computer_tools import shell as shell_tools
|
|
|
|
calls = []
|
|
|
|
class FakeShell:
|
|
async def exec(self, command, cwd=None, background=False, env=None):
|
|
calls.append({"command": command, "background": background})
|
|
return {"success": True, "stdout": "", "stderr": "", "exit_code": 0}
|
|
|
|
class FakeBooter:
|
|
shell = FakeShell()
|
|
|
|
FakeWrapper = _make_fake_wrapper_class()
|
|
|
|
async def fake_get_booter(context, session_id):
|
|
return FakeBooter()
|
|
|
|
monkeypatch.setattr(shell_tools, "get_booter", fake_get_booter)
|
|
|
|
result = await ExecuteShellTool().call(
|
|
FakeWrapper(), command="chromium https://example.com"
|
|
)
|
|
|
|
assert isinstance(result, str)
|
|
assert json.loads(result)["success"] is True
|
|
assert calls == [{"command": "chromium https://example.com", "background": False}]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_shell_uses_fresh_default_env_per_call(monkeypatch):
|
|
from astrbot.core.tools.computer_tools import shell as shell_tools
|
|
|
|
calls = []
|
|
|
|
class FakeShell:
|
|
async def exec(self, command, cwd=None, background=False, env=None):
|
|
assert env is not None
|
|
env["MUTATED_BY_FAKE_SHELL"] = command
|
|
calls.append(env.copy())
|
|
return {"success": True, "stdout": "", "stderr": "", "exit_code": 0}
|
|
|
|
class FakeBooter:
|
|
shell = FakeShell()
|
|
|
|
FakeWrapper = _make_fake_wrapper_class()
|
|
|
|
async def fake_get_booter(context, session_id):
|
|
return FakeBooter()
|
|
|
|
monkeypatch.setattr(shell_tools, "get_booter", fake_get_booter)
|
|
tool = ExecuteShellTool()
|
|
|
|
await tool.call(FakeWrapper(), command="first")
|
|
await tool.call(FakeWrapper(), command="second")
|
|
|
|
assert calls[0]["MUTATED_BY_FAKE_SHELL"] == "first"
|
|
assert calls[1] == {"MUTATED_BY_FAKE_SHELL": "second"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_shell_copies_user_env_before_execution(monkeypatch):
|
|
from astrbot.core.tools.computer_tools import shell as shell_tools
|
|
|
|
calls = []
|
|
|
|
class FakeShell:
|
|
async def exec(self, command, cwd=None, background=False, env=None):
|
|
assert env is not None
|
|
env["MUTATED_BY_FAKE_SHELL"] = command
|
|
calls.append(env.copy())
|
|
return {"success": True, "stdout": "", "stderr": "", "exit_code": 0}
|
|
|
|
class FakeBooter:
|
|
shell = FakeShell()
|
|
|
|
FakeWrapper = _make_fake_wrapper_class()
|
|
|
|
async def fake_get_booter(context, session_id):
|
|
return FakeBooter()
|
|
|
|
monkeypatch.setattr(shell_tools, "get_booter", fake_get_booter)
|
|
original_env = {"FOO": "bar"}
|
|
|
|
await ExecuteShellTool().call(FakeWrapper(), command="first", env=original_env)
|
|
|
|
assert original_env == {"FOO": "bar"}
|
|
assert calls == [{"FOO": "bar", "MUTATED_BY_FAKE_SHELL": "first"}]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_shell_passes_background_flag_directly(monkeypatch):
|
|
"""In the new architecture, background flag is passed directly to shell exec."""
|
|
from astrbot.core.tools.computer_tools import shell as shell_tools
|
|
|
|
calls = []
|
|
|
|
class FakeShell:
|
|
async def exec(self, command, cwd=None, background=False, env=None):
|
|
calls.append({"command": command, "background": background})
|
|
return {"success": True, "stdout": "", "stderr": "", "exit_code": 0}
|
|
|
|
class FakeBooter:
|
|
shell = FakeShell()
|
|
|
|
FakeWrapper = _make_fake_wrapper_class()
|
|
|
|
async def fake_get_booter(context, session_id):
|
|
return FakeBooter()
|
|
|
|
monkeypatch.setattr(shell_tools, "get_booter", fake_get_booter)
|
|
|
|
command = "nohup firefox >/tmp/astrbot-firefox.log 2>&1 &"
|
|
result = await ExecuteShellTool().call(
|
|
FakeWrapper(), command=command, background=True
|
|
)
|
|
|
|
assert isinstance(result, str)
|
|
assert json.loads(result)["success"] is True
|
|
assert calls == [{"command": command, "background": True}]
|
|
|
|
command2 = "firefox & # already detached"
|
|
result2 = await ExecuteShellTool().call(
|
|
FakeWrapper(), command=command2, background=True
|
|
)
|
|
|
|
assert isinstance(result2, str)
|
|
assert json.loads(result2)["success"] is True
|
|
assert calls[1] == {"command": command2, "background": True}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_shell_reports_exception_type(monkeypatch):
|
|
"""Error message uses e!s formatting (may omit class name if __str__ is blank)."""
|
|
from astrbot.core.tools.computer_tools import shell as shell_tools
|
|
|
|
class FakeShell:
|
|
async def exec(self, command, cwd=None, background=False, env=None):
|
|
raise ValueError("custom error")
|
|
|
|
class FakeBooter:
|
|
shell = FakeShell()
|
|
|
|
FakeWrapper = _make_fake_wrapper_class()
|
|
|
|
async def fake_get_booter(context, session_id):
|
|
return FakeBooter()
|
|
|
|
monkeypatch.setattr(shell_tools, "get_booter", fake_get_booter)
|
|
|
|
result = await ExecuteShellTool().call(FakeWrapper(), command="firefox")
|
|
|
|
assert result == "Error executing command: custom error"
|
|
|
|
|
|
def test_tavily_tools_are_registered_as_builtin_tools():
|
|
manager = FunctionToolManager()
|
|
|
|
search_tool = manager.get_builtin_tool(TavilyWebSearchTool)
|
|
extract_tool = manager.get_builtin_tool(TavilyExtractWebPageTool)
|
|
|
|
assert tool.name == "astrbot_execute_shell"
|
|
assert tool.parameters["properties"]["background"]["default"] is False
|
|
assert manager.is_builtin_tool("astrbot_execute_shell") is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_shell_defaults_to_foreground(monkeypatch):
|
|
from astrbot.core.tools.computer_tools import shell as shell_tools
|
|
|
|
calls = []
|
|
|
|
class FakeShell:
|
|
async def exec(
|
|
self, command, cwd=None, background=False, env=None, timeout=None
|
|
):
|
|
calls.append({"command": command, "background": background})
|
|
return {"success": True, "stdout": "", "stderr": "", "exit_code": 0}
|
|
|
|
class FakeBooter:
|
|
shell = FakeShell()
|
|
|
|
class FakeConfig:
|
|
def get_config(self, umo):
|
|
return {"provider_settings": {"computer_use_runtime": "sandbox"}}
|
|
|
|
class FakeEvent:
|
|
unified_msg_origin = "umo"
|
|
role = "admin"
|
|
|
|
class FakeAstrContext:
|
|
context = FakeConfig()
|
|
event = FakeEvent()
|
|
|
|
class FakeWrapper:
|
|
context = FakeAstrContext()
|
|
|
|
async def fake_get_booter(context, session_id):
|
|
return FakeBooter()
|
|
|
|
monkeypatch.setattr(shell_tools, "get_booter", fake_get_booter)
|
|
|
|
result = await ExecuteShellTool().call(
|
|
FakeWrapper(), command="chromium https://example.com"
|
|
)
|
|
|
|
assert json.loads(result)["success"] is True
|
|
assert calls == [{"command": "chromium https://example.com", "background": False}]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_shell_uses_fresh_default_env_per_call(monkeypatch):
|
|
from astrbot.core.tools.computer_tools import shell as shell_tools
|
|
|
|
calls = []
|
|
|
|
class FakeShell:
|
|
async def exec(
|
|
self, command, cwd=None, background=False, env=None, timeout=None
|
|
):
|
|
env["MUTATED_BY_FAKE_SHELL"] = command
|
|
calls.append(env)
|
|
return {"success": True, "stdout": "", "stderr": "", "exit_code": 0}
|
|
|
|
class FakeBooter:
|
|
shell = FakeShell()
|
|
|
|
class FakeConfig:
|
|
def get_config(self, umo):
|
|
return {"provider_settings": {"computer_use_runtime": "sandbox"}}
|
|
|
|
class FakeEvent:
|
|
unified_msg_origin = "umo"
|
|
role = "admin"
|
|
|
|
class FakeAstrContext:
|
|
context = FakeConfig()
|
|
event = FakeEvent()
|
|
|
|
class FakeWrapper:
|
|
context = FakeAstrContext()
|
|
|
|
async def fake_get_booter(context, session_id):
|
|
return FakeBooter()
|
|
|
|
monkeypatch.setattr(shell_tools, "get_booter", fake_get_booter)
|
|
tool = ExecuteShellTool()
|
|
|
|
await tool.call(FakeWrapper(), command="first")
|
|
await tool.call(FakeWrapper(), command="second")
|
|
|
|
assert calls[0] is not calls[1]
|
|
assert calls[0]["MUTATED_BY_FAKE_SHELL"] == "first"
|
|
assert calls[1] == {"MUTATED_BY_FAKE_SHELL": "second"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_shell_copies_user_env_before_execution(monkeypatch):
|
|
from astrbot.core.tools.computer_tools import shell as shell_tools
|
|
|
|
calls = []
|
|
|
|
class FakeShell:
|
|
async def exec(
|
|
self, command, cwd=None, background=False, env=None, timeout=None
|
|
):
|
|
env["MUTATED_BY_FAKE_SHELL"] = command
|
|
calls.append(env)
|
|
return {"success": True, "stdout": "", "stderr": "", "exit_code": 0}
|
|
|
|
class FakeBooter:
|
|
shell = FakeShell()
|
|
|
|
class FakeConfig:
|
|
def get_config(self, umo):
|
|
return {"provider_settings": {"computer_use_runtime": "sandbox"}}
|
|
|
|
class FakeEvent:
|
|
unified_msg_origin = "umo"
|
|
role = "admin"
|
|
|
|
class FakeAstrContext:
|
|
context = FakeConfig()
|
|
event = FakeEvent()
|
|
|
|
class FakeWrapper:
|
|
context = FakeAstrContext()
|
|
|
|
async def fake_get_booter(context, session_id):
|
|
return FakeBooter()
|
|
|
|
monkeypatch.setattr(shell_tools, "get_booter", fake_get_booter)
|
|
original_env = {"FOO": "bar"}
|
|
|
|
await ExecuteShellTool().call(FakeWrapper(), command="first", env=original_env)
|
|
|
|
assert original_env == {"FOO": "bar"}
|
|
assert calls == [{"FOO": "bar", "MUTATED_BY_FAKE_SHELL": "first"}]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_shell_avoids_double_background_for_detached_commands(
|
|
monkeypatch,
|
|
):
|
|
from astrbot.core.tools.computer_tools import shell as shell_tools
|
|
|
|
calls = []
|
|
|
|
class FakeShell:
|
|
async def exec(
|
|
self, command, cwd=None, background=False, env=None, timeout=None
|
|
):
|
|
calls.append({"command": command, "background": background})
|
|
return {"success": True, "stdout": "", "stderr": "", "exit_code": 0}
|
|
|
|
class FakeBooter:
|
|
shell = FakeShell()
|
|
|
|
class FakeConfig:
|
|
def get_config(self, umo):
|
|
return {"provider_settings": {"computer_use_runtime": "sandbox"}}
|
|
|
|
class FakeEvent:
|
|
unified_msg_origin = "umo"
|
|
role = "admin"
|
|
|
|
class FakeAstrContext:
|
|
context = FakeConfig()
|
|
event = FakeEvent()
|
|
|
|
class FakeWrapper:
|
|
context = FakeAstrContext()
|
|
|
|
async def fake_get_booter(context, session_id):
|
|
return FakeBooter()
|
|
|
|
monkeypatch.setattr(shell_tools, "get_booter", fake_get_booter)
|
|
|
|
command = "nohup firefox >/tmp/astrbot-firefox.log 2>&1 &"
|
|
result = await ExecuteShellTool().call(
|
|
FakeWrapper(), command=command, background=True
|
|
)
|
|
|
|
assert json.loads(result)["success"] is True
|
|
assert calls == [{"command": command, "background": False}]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_shell_recognizes_commented_background_command(monkeypatch):
|
|
from astrbot.core.tools.computer_tools import shell as shell_tools
|
|
|
|
calls = []
|
|
|
|
class FakeShell:
|
|
async def exec(
|
|
self, command, cwd=None, background=False, env=None, timeout=None
|
|
):
|
|
calls.append({"command": command, "background": background})
|
|
return {"success": True, "stdout": "", "stderr": "", "exit_code": 0}
|
|
|
|
class FakeBooter:
|
|
shell = FakeShell()
|
|
|
|
class FakeConfig:
|
|
def get_config(self, umo):
|
|
return {"provider_settings": {"computer_use_runtime": "sandbox"}}
|
|
|
|
class FakeEvent:
|
|
unified_msg_origin = "umo"
|
|
role = "admin"
|
|
|
|
class FakeAstrContext:
|
|
context = FakeConfig()
|
|
event = FakeEvent()
|
|
|
|
class FakeWrapper:
|
|
context = FakeAstrContext()
|
|
|
|
async def fake_get_booter(context, session_id):
|
|
return FakeBooter()
|
|
|
|
monkeypatch.setattr(shell_tools, "get_booter", fake_get_booter)
|
|
|
|
command = "firefox & # already detached"
|
|
result = await ExecuteShellTool().call(
|
|
FakeWrapper(), command=command, background=True
|
|
)
|
|
|
|
assert json.loads(result)["success"] is True
|
|
assert calls == [{"command": command, "background": False}]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("command", "expected"),
|
|
[
|
|
("echo '#'", False),
|
|
("echo '&'", False),
|
|
("echo foo#bar &", True),
|
|
("echo 'unterminated", False),
|
|
("firefox & # already detached", True),
|
|
("nohup firefox >/tmp/astrbot-firefox.log 2>&1 &", True),
|
|
("firefox", False),
|
|
],
|
|
)
|
|
def test_is_self_detached_command_handles_quotes_and_comments(command, expected):
|
|
from astrbot.core.tools.computer_tools.shell import _is_self_detached_command
|
|
|
|
assert _is_self_detached_command(command) is expected
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_execute_shell_reports_blank_exception_type(monkeypatch):
|
|
from astrbot.core.tools.computer_tools import shell as shell_tools
|
|
|
|
class BlankError(Exception):
|
|
def __str__(self):
|
|
return ""
|
|
|
|
class FakeShell:
|
|
async def exec(
|
|
self, command, cwd=None, background=False, env=None, timeout=None
|
|
):
|
|
raise BlankError()
|
|
|
|
class FakeBooter:
|
|
shell = FakeShell()
|
|
|
|
class FakeConfig:
|
|
def get_config(self, umo):
|
|
return {"provider_settings": {"computer_use_runtime": "sandbox"}}
|
|
|
|
class FakeEvent:
|
|
unified_msg_origin = "umo"
|
|
role = "admin"
|
|
|
|
class FakeAstrContext:
|
|
context = FakeConfig()
|
|
event = FakeEvent()
|
|
|
|
class FakeWrapper:
|
|
context = FakeAstrContext()
|
|
|
|
async def fake_get_booter(context, session_id):
|
|
return FakeBooter()
|
|
|
|
monkeypatch.setattr(shell_tools, "get_booter", fake_get_booter)
|
|
|
|
result = await ExecuteShellTool().call(FakeWrapper(), command="firefox")
|
|
|
|
assert result == "Error executing command: BlankError"
|
|
|
|
|
|
def test_firecrawl_tools_are_registered_as_builtin_tools():
|
|
manager = FunctionToolManager()
|
|
|
|
search_tool = manager.get_builtin_tool(FirecrawlWebSearchTool)
|
|
extract_tool = manager.get_builtin_tool(FirecrawlExtractWebPageTool)
|
|
|
|
assert search_tool.name == "web_search_firecrawl"
|
|
assert extract_tool.name == "firecrawl_extract_web_page"
|
|
assert manager.is_builtin_tool("web_search_firecrawl") is True
|
|
assert manager.is_builtin_tool("firecrawl_extract_web_page") is True
|