refactor: improve astrbot builtin tool management#7418
Conversation
- Introduced a new registry for builtin tools to streamline their management. - Added `SendMessageToUserTool`, `KnowledgeBaseQueryTool`, and various web search tools as builtin tools. - Updated `FunctionToolManager` to cache and retrieve builtin tools efficiently. - Modified `CronJobManager` to utilize the new `SendMessageToUserTool`. - Enhanced the dashboard to display readonly status for builtin tools and prevent toggling their state. - Added tests for builtin tool injection and retrieval to ensure proper functionality.
There was a problem hiding this comment.
Hey - I've found 3 issues, and left some high level feedback:
- In
FunctionToolManager.get_func, theif getattr(builtin_tool, "active", True): return builtin_toolfollowed byreturn builtin_toolalways returns the tool regardless ofactive, so either drop the condition or make the inactive case returnNone(or otherwise behave differently) to reflect the intent.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- In `FunctionToolManager.get_func`, the `if getattr(builtin_tool, "active", True): return builtin_tool` followed by `return builtin_tool` always returns the tool regardless of `active`, so either drop the condition or make the inactive case return `None` (or otherwise behave differently) to reflect the intent.
## Individual Comments
### Comment 1
<location path="astrbot/core/tools/message_tools.py" line_range="76-85" />
<code_context>
- }
- )
-
- async def _resolve_path_from_sandbox(
- self, context: ContextWrapper[AstrAgentContext], path: str
- ) -> tuple[str, bool]:
- """
- If the path exists locally, return it directly.
- Otherwise, check if it exists in the sandbox and download it.
-
- bool: indicates whether the file was downloaded from sandbox.
- """
- if os.path.exists(path):
- return path, False
-
- # Try to check if the file exists in the sandbox
- try:
- sb = await get_booter(
- context.context.context,
- context.context.event.unified_msg_origin,
- )
- # Use shell to check if the file exists in sandbox
- result = await sb.shell.exec(f"test -f {path} && echo '_&exists_'")
- if "_&exists_" in json.dumps(result):
- # Download the file from sandbox
</code_context>
<issue_to_address>
**🚨 issue (security):** Use safer shell invocation or quoting when interpolating file paths into commands.
This uses `sb.shell.exec(f"test -f {path} && echo '_&exists_'")` with an unquoted `path`. If `path` can contain spaces or shell metacharacters, the check may break and it also introduces a command-injection risk. Prefer either passing `path` as an argument to `exec` (no shell) if supported, or quoting it first (e.g., via `shlex.quote`). Since this helper is now more general, it’s a good place to harden this behavior.
</issue_to_address>
### Comment 2
<location path="astrbot/dashboard/routes/tools.py" line_range="439-444" />
<code_context>
tools_dict = []
for tool in tools:
- if isinstance(tool, MCPTool):
+ readonly = False
+ if self.tool_mgr.is_builtin_tool(tool.name):
+ origin = "builtin"
+ origin_name = "AstrBot Core"
+ readonly = True
+ elif isinstance(tool, MCPTool):
origin = "mcp"
origin_name = tool.mcp_server_name
</code_context>
<issue_to_address>
**issue (bug_risk):** Name-based builtin detection can misclassify tools if a plugin/MCP tool uses a reserved builtin name.
`origin` and `readonly` are derived only from `self.tool_mgr.is_builtin_tool(tool.name)`, so any plugin/MCP tool that happens to share a name with a builtin will be treated as `origin='builtin'` and `readonly=True`, and `toggle_tool` will then block enabling/disabling it. If builtin names are not strictly reserved, consider basing this on the actual tool implementation (e.g., class or source module) rather than just the name.
</issue_to_address>
### Comment 3
<location path="tests/unit/test_func_tool_manager.py" line_range="1-10" />
<code_context>
from astrbot.core import sp
from astrbot.core.agent.mcp_client import MCPClient, MCPTool
from astrbot.core.agent.tool import FunctionTool, ToolSet
</code_context>
<issue_to_address>
**suggestion (testing):** Broaden `FunctionToolManager` tests to cover error paths and builtin-tool discovery helpers
The current test confirms caching behavior and name resolution, but the new responsibilities merit broader coverage. Consider adding tests that:
- Assert `get_builtin_tool` raises `KeyError` for unknown string names and `TypeError` for invalid argument types.
- Assert `get_builtin_tool` raises `KeyError` when given an unregistered `FunctionTool` subclass (exercising the `get_builtin_tool_name` path).
- Verify `iter_builtin_tools` yields at least the expected builtins (e.g., includes `SendMessageToUserTool`) and that returned instances match those from `get_builtin_tool`.
- Verify `is_builtin_tool` is `True` for known builtin names and `False` for non-builtin tools.
These will better exercise both happy-path and error-path behavior of the builtin registry integration.
Suggested implementation:
```python
import pytest
from astrbot.core import sp
from astrbot.core.agent.tool import FunctionTool
from astrbot.core.provider.func_tool_manager import FunctionToolManager
from astrbot.core.tools.message_tools import SendMessageToUserTool
```
```python
class UnregisteredFunctionTool(FunctionTool):
"""Dummy FunctionTool subclass that is not registered as a builtin."""
name = "unregistered_function_tool"
def test_get_builtin_tool_by_class_returns_cached_instance():
manager = FunctionToolManager()
tool_by_class = manager.get_builtin_tool(SendMessageToUserTool)
tool_by_name = manager.get_builtin_tool("send_message_to_user")
assert tool_by_class is tool_by_name
assert manager.get_func("send_message_to_user") is tool_by_class
assert tool_by_class.name == "send_message_to_user"
def test_get_builtin_tool_raises_key_error_for_unknown_name():
manager = FunctionToolManager()
with pytest.raises(KeyError):
manager.get_builtin_tool("non_existent_tool_name")
def test_get_builtin_tool_raises_type_error_for_invalid_argument_type():
manager = FunctionToolManager()
with pytest.raises(TypeError):
manager.get_builtin_tool(123) # type: ignore[arg-type]
with pytest.raises(TypeError):
manager.get_builtin_tool(object()) # type: ignore[arg-type]
def test_get_builtin_tool_raises_key_error_for_unregistered_subclass():
manager = FunctionToolManager()
with pytest.raises(KeyError):
manager.get_builtin_tool(UnregisteredFunctionTool)
def test_iter_builtin_tools_includes_send_message_to_user_and_matches_get_builtin_tool():
manager = FunctionToolManager()
all_tools = list(manager.iter_builtin_tools())
send_message_tools = [
tool for tool in all_tools if isinstance(tool, SendMessageToUserTool)
]
# At least one SendMessageToUserTool is present in the builtin registry.
assert send_message_tools
tool_from_get = manager.get_builtin_tool(SendMessageToUserTool)
# The instance returned from get_builtin_tool should be the same as one
# yielded by iter_builtin_tools.
assert any(tool is tool_from_get for tool in send_message_tools)
def test_is_builtin_tool_for_known_and_unknown_tools():
manager = FunctionToolManager()
assert manager.is_builtin_tool("send_message_to_user") is True
assert manager.is_builtin_tool(SendMessageToUserTool) is True
assert manager.is_builtin_tool("non_existent_tool_name") is False
assert manager.is_builtin_tool(UnregisteredFunctionTool) is False
```
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
| readonly = False | ||
| if self.tool_mgr.is_builtin_tool(tool.name): | ||
| origin = "builtin" | ||
| origin_name = "AstrBot Core" | ||
| readonly = True | ||
| elif isinstance(tool, MCPTool): |
There was a problem hiding this comment.
issue (bug_risk): Name-based builtin detection can misclassify tools if a plugin/MCP tool uses a reserved builtin name.
origin and readonly are derived only from self.tool_mgr.is_builtin_tool(tool.name), so any plugin/MCP tool that happens to share a name with a builtin will be treated as origin='builtin' and readonly=True, and toggle_tool will then block enabling/disabling it. If builtin names are not strictly reserved, consider basing this on the actual tool implementation (e.g., class or source module) rather than just the name.
| from astrbot.core import sp | ||
| from astrbot.core.provider.func_tool_manager import FunctionToolManager | ||
| from astrbot.core.tools.message_tools import SendMessageToUserTool | ||
|
|
||
|
|
||
| def test_get_builtin_tool_by_class_returns_cached_instance(): | ||
| manager = FunctionToolManager() | ||
|
|
||
| tool_by_class = manager.get_builtin_tool(SendMessageToUserTool) | ||
| tool_by_name = manager.get_builtin_tool("send_message_to_user") |
There was a problem hiding this comment.
suggestion (testing): Broaden FunctionToolManager tests to cover error paths and builtin-tool discovery helpers
The current test confirms caching behavior and name resolution, but the new responsibilities merit broader coverage. Consider adding tests that:
- Assert
get_builtin_toolraisesKeyErrorfor unknown string names andTypeErrorfor invalid argument types. - Assert
get_builtin_toolraisesKeyErrorwhen given an unregisteredFunctionToolsubclass (exercising theget_builtin_tool_namepath). - Verify
iter_builtin_toolsyields at least the expected builtins (e.g., includesSendMessageToUserTool) and that returned instances match those fromget_builtin_tool. - Verify
is_builtin_toolisTruefor known builtin names andFalsefor non-builtin tools.
These will better exercise both happy-path and error-path behavior of the builtin registry integration.
Suggested implementation:
import pytest
from astrbot.core import sp
from astrbot.core.agent.tool import FunctionTool
from astrbot.core.provider.func_tool_manager import FunctionToolManager
from astrbot.core.tools.message_tools import SendMessageToUserToolclass UnregisteredFunctionTool(FunctionTool):
"""Dummy FunctionTool subclass that is not registered as a builtin."""
name = "unregistered_function_tool"
def test_get_builtin_tool_by_class_returns_cached_instance():
manager = FunctionToolManager()
tool_by_class = manager.get_builtin_tool(SendMessageToUserTool)
tool_by_name = manager.get_builtin_tool("send_message_to_user")
assert tool_by_class is tool_by_name
assert manager.get_func("send_message_to_user") is tool_by_class
assert tool_by_class.name == "send_message_to_user"
def test_get_builtin_tool_raises_key_error_for_unknown_name():
manager = FunctionToolManager()
with pytest.raises(KeyError):
manager.get_builtin_tool("non_existent_tool_name")
def test_get_builtin_tool_raises_type_error_for_invalid_argument_type():
manager = FunctionToolManager()
with pytest.raises(TypeError):
manager.get_builtin_tool(123) # type: ignore[arg-type]
with pytest.raises(TypeError):
manager.get_builtin_tool(object()) # type: ignore[arg-type]
def test_get_builtin_tool_raises_key_error_for_unregistered_subclass():
manager = FunctionToolManager()
with pytest.raises(KeyError):
manager.get_builtin_tool(UnregisteredFunctionTool)
def test_iter_builtin_tools_includes_send_message_to_user_and_matches_get_builtin_tool():
manager = FunctionToolManager()
all_tools = list(manager.iter_builtin_tools())
send_message_tools = [
tool for tool in all_tools if isinstance(tool, SendMessageToUserTool)
]
# At least one SendMessageToUserTool is present in the builtin registry.
assert send_message_tools
tool_from_get = manager.get_builtin_tool(SendMessageToUserTool)
# The instance returned from get_builtin_tool should be the same as one
# yielded by iter_builtin_tools.
assert any(tool is tool_from_get for tool in send_message_tools)
def test_is_builtin_tool_for_known_and_unknown_tools():
manager = FunctionToolManager()
assert manager.is_builtin_tool("send_message_to_user") is True
assert manager.is_builtin_tool(SendMessageToUserTool) is True
assert manager.is_builtin_tool("non_existent_tool_name") is False
assert manager.is_builtin_tool(UnregisteredFunctionTool) is FalseThere was a problem hiding this comment.
Code Review
This pull request introduces a centralized registry and management system for builtin tools, refactoring existing tools (Knowledge Base, Messaging, Cron, and Web Search) into a modular structure under astrbot/core/tools/. The FunctionToolManager was updated to handle the lifecycle of these tools, and the dashboard now correctly identifies and displays them as read-only. Key feedback includes a critical security concern regarding command injection in the sandbox file resolution logic, a recommendation to track temporary files for cleanup to prevent resource leaks, and a suggestion to simplify redundant logic in the tool manager's retrieval method.
* feat: Refactor astrbot builtin tool management - Introduced a new registry for builtin tools to streamline their management. - Added `SendMessageToUserTool`, `KnowledgeBaseQueryTool`, and various web search tools as builtin tools. - Updated `FunctionToolManager` to cache and retrieve builtin tools efficiently. - Modified `CronJobManager` to utilize the new `SendMessageToUserTool`. - Enhanced the dashboard to display readonly status for builtin tools and prevent toggling their state. - Added tests for builtin tool injection and retrieval to ensure proper functionality. * fix: escape file path in shell command to prevent injection vulnerabilities
SendMessageToUserTool,KnowledgeBaseQueryTool, and various web search tools as builtin tools.FunctionToolManagerto cache and retrieve builtin tools efficiently.CronJobManagerto utilize the newSendMessageToUserTool.Modifications / 改动点
Screenshots or Test Results / 运行截图或测试结果
Checklist / 检查清单
😊 If there are new features added in the PR, I have discussed it with the authors through issues/emails, etc.
/ 如果 PR 中有新加入的功能,已经通过 Issue / 邮件等方式和作者讨论过。
👀 My changes have been well-tested, and "Verification Steps" and "Screenshots" have been provided above.
/ 我的更改经过了良好的测试,并已在上方提供了“验证步骤”和“运行截图”。
🤓 I have ensured that no new dependencies are introduced, OR if new dependencies are introduced, they have been added to the appropriate locations in
requirements.txtandpyproject.toml./ 我确保没有引入新依赖库,或者引入了新依赖库的同时将其添加到
requirements.txt和pyproject.toml文件相应位置。😮 My changes do not introduce malicious code.
/ 我的更改没有引入恶意代码。
Summary by Sourcery
Introduce a centralized registry and lifecycle for AstrBot builtin tools and adopt it across core agents, cron, and web search features.
New Features:
Enhancements:
Tests: