diff --git a/astrbot/core/star/context.py b/astrbot/core/star/context.py index 593bad936..4cf0e42a0 100644 --- a/astrbot/core/star/context.py +++ b/astrbot/core/star/context.py @@ -51,6 +51,67 @@ if TYPE_CHECKING: WebApiHandler = Callable[..., Awaitable[Any]] RegisteredWebApi = tuple[str, WebApiHandler, list[str], str] +_PLUGIN_MODULE_FLAGS = {"builtin_stars", "plugins"} + + +def _split_module_path(module_path: Any) -> list[str]: + if not isinstance(module_path, str) or not module_path: + return [] + return module_path.split(".") + + +def _plugin_root_from_module_parts(parts: list[str]) -> tuple[str, str] | None: + for index, part in enumerate(parts): + if part in _PLUGIN_MODULE_FLAGS and index + 1 < len(parts): + return part, parts[index + 1] + return None + + +def _plugin_root_from_metadata(metadata: StarMetadata) -> str | None: + if metadata.root_dir_name: + return metadata.root_dir_name + + root_info = _plugin_root_from_module_parts(_split_module_path(metadata.module_path)) + return root_info[1] if root_info else None + + +def _registered_plugin_module_path(root_dir_name: str, flag: str | None) -> str | None: + for metadata in reversed(star_registry): + if not metadata.module_path: + continue + if _plugin_root_from_metadata(metadata) != root_dir_name: + continue + if flag and flag not in _split_module_path(metadata.module_path): + continue + return metadata.module_path + return None + + +def _legacy_plugin_module_path(parts: list[str]) -> str: + resolved_parts = [] + for index, part in enumerate(parts): + resolved_parts.append(part) + if part in _PLUGIN_MODULE_FLAGS and index + 1 < len(parts): + resolved_parts.append(parts[index + 1]) + resolved_parts.append("main") + break + return ".".join(resolved_parts) + + +def _resolve_tool_handler_module_path(tool: FunctionTool) -> str: + module_path = getattr(tool, "__module__", None) + module_parts = _split_module_path(module_path) + if not module_parts: + return module_path if isinstance(module_path, str) else "" + + root_info = _plugin_root_from_module_parts(module_parts) + if root_info: + flag, root_dir_name = root_info + registered_module_path = _registered_plugin_module_path(root_dir_name, flag) + return registered_module_path or _legacy_plugin_module_path(module_parts) + + registered_module_path = _registered_plugin_module_path(module_parts[0], "plugins") + return registered_module_path or ".".join(module_parts) class PlatformManagerProtocol(Protocol): @@ -490,16 +551,7 @@ class Context: module_path = "" for tool in tools: if not module_path: - _parts = [] - module_part = tool.__module__.split(".") - flags = ["builtin_stars", "plugins"] - for i, part in enumerate(module_part): - _parts.append(part) - if part in flags and i + 1 < len(module_part): - _parts.append(module_part[i + 1]) - _parts.append("main") - break - tool.handler_module_path = ".".join(_parts) + tool.handler_module_path = _resolve_tool_handler_module_path(tool) module_path = tool.handler_module_path else: tool.handler_module_path = module_path diff --git a/tests/unit/test_star_context.py b/tests/unit/test_star_context.py new file mode 100644 index 000000000..0979aaa0b --- /dev/null +++ b/tests/unit/test_star_context.py @@ -0,0 +1,106 @@ +from types import SimpleNamespace + +import pytest + +from astrbot.core.agent.tool import FunctionTool +from astrbot.core.provider.func_tool_manager import FunctionToolManager +from astrbot.core.star.context import Context +from astrbot.core.star.star import StarMetadata, star_registry + + +@pytest.fixture(autouse=True) +def restore_star_registry(): + original_registry = list(star_registry) + star_registry.clear() + try: + yield + finally: + star_registry[:] = original_registry + + +def make_context() -> Context: + context = Context.__new__(Context) + context.provider_manager = SimpleNamespace(llm_tools=FunctionToolManager()) + return context + + +def make_tool(name: str, module_path: str) -> FunctionTool: + tool = FunctionTool( + name=name, + description="test tool", + parameters={"type": "object", "properties": {}}, + ) + tool.__module__ = module_path + return tool + + +def test_add_llm_tools_resolves_subdirectory_plugin_without_name_prefix(): + star_registry.append( + StarMetadata( + name="Custom Plugin", + root_dir_name="custom_plugin", + module_path="data.plugins.custom_plugin.main", + ) + ) + context = make_context() + tool = make_tool("search", "custom_plugin.tools.search") + + context.add_llm_tools(tool) + + assert tool.handler_module_path == "data.plugins.custom_plugin.main" + + +def test_add_llm_tools_uses_registered_non_main_plugin_entrypoint(): + star_registry.append( + StarMetadata( + name="Custom Plugin", + module_path="data.plugins.custom_plugin.custom_plugin", + ) + ) + context = make_context() + tool = make_tool("search", "custom_plugin.tools.search") + + context.add_llm_tools(tool) + + assert tool.handler_module_path == "data.plugins.custom_plugin.custom_plugin" + + +def test_add_llm_tools_resolves_prefixed_subdirectory_tool_from_registry(): + star_registry.append( + StarMetadata( + name="Custom Plugin", + root_dir_name="custom_plugin", + module_path="data.plugins.custom_plugin.custom_plugin", + ) + ) + context = make_context() + tool = make_tool("search", "data.plugins.custom_plugin.tools.search") + + context.add_llm_tools(tool) + + assert tool.handler_module_path == "data.plugins.custom_plugin.custom_plugin" + + +def test_add_llm_tools_does_not_treat_unknown_module_as_plugin(): + star_registry.append( + StarMetadata( + name="Custom Plugin", + root_dir_name="custom_plugin", + module_path="data.plugins.custom_plugin.main", + ) + ) + context = make_context() + tool = make_tool("search", "external_package.tools.search") + + context.add_llm_tools(tool) + + assert tool.handler_module_path == "external_package.tools.search" + + +def test_add_llm_tools_handles_empty_tool_module_path(): + context = make_context() + tool = make_tool("search", "") + + context.add_llm_tools(tool) + + assert tool.handler_module_path == ""