diff --git a/services/sandbox/install_tool_shims.py b/services/sandbox/install_tool_shims.py index 9db44ea39..06a140fb2 100644 --- a/services/sandbox/install_tool_shims.py +++ b/services/sandbox/install_tool_shims.py @@ -140,6 +140,20 @@ def _clear_published_tools(tool_dir: Path) -> None: _remove_path(child) +def _tool_project_name(package_dir: Path) -> str: + pyproject = package_dir / "pyproject.toml" + try: + data = tomllib.loads(pyproject.read_text()) + except (OSError, tomllib.TOMLDecodeError) as exc: + print(f"warning: failed to read {pyproject}: {exc}", file=sys.stderr) + return "" + project = data.get("project") or {} + if not isinstance(project, dict): + print(f"warning: invalid [project] table in {pyproject}", file=sys.stderr) + return "" + return str(project.get("name") or "") + + def _copy_published_tools(tool_dir: Path, published: Path) -> None: if not published.is_dir(): raise RuntimeError(f"refreshed tools subdir does not exist: {published}") @@ -148,16 +162,19 @@ def _copy_published_tools(tool_dir: Path, published: Path) -> None: blocklist = _tool_blocklist() existing = {package_dir.name: package_dir for package_dir in _tool_package_dirs(tool_dir)} for package_dir in _tool_package_dirs(published): - tool_name = package_dir.name - if allowlist is not None and tool_name not in allowlist: + package_dir_name = package_dir.name + filter_names = {package_dir_name} + if allowlist is not None or blocklist: + filter_names.add(_tool_project_name(package_dir)) + if allowlist is not None and allowlist.isdisjoint(filter_names): # Not in TOOL_ALLOWLIST -> don't install; keeps the agent's catalog # to configured tools (no phantom, credential-less tools). continue - if tool_name in blocklist: + if not blocklist.isdisjoint(filter_names): continue - if tool_name in existing: + if package_dir_name in existing: print( - f"skipping duplicate tool {tool_name}: {package_dir} conflicts with {existing[tool_name]}", + f"skipping duplicate tool {package_dir_name}: {package_dir} conflicts with {existing[package_dir_name]}", file=sys.stderr, ) continue @@ -167,7 +184,7 @@ def _copy_published_tools(tool_dir: Path, published: Path) -> None: _remove_path(target) target.parent.mkdir(parents=True, exist_ok=True) shutil.copytree(package_dir, target, symlinks=True) - existing[tool_name] = target + existing[package_dir_name] = target def _tool_package_dirs(published: Path) -> list[Path]: diff --git a/services/sandbox/test_install_tool_shims.py b/services/sandbox/test_install_tool_shims.py index 672d0027d..3a15324ba 100644 --- a/services/sandbox/test_install_tool_shims.py +++ b/services/sandbox/test_install_tool_shims.py @@ -1,16 +1,24 @@ from __future__ import annotations import contextlib +import importlib.util import io import json import os +from pathlib import Path import subprocess import tempfile import unittest -from pathlib import Path from unittest import mock -import install_tool_shims +INSTALL_TOOL_SHIMS_PATH = Path(__file__).resolve().parent / "install_tool_shims.py" +INSTALL_TOOL_SHIMS_SPEC = importlib.util.spec_from_file_location( + "install_tool_shims", INSTALL_TOOL_SHIMS_PATH +) +if INSTALL_TOOL_SHIMS_SPEC is None or INSTALL_TOOL_SHIMS_SPEC.loader is None: + raise RuntimeError(f"failed to load {INSTALL_TOOL_SHIMS_PATH}") +install_tool_shims = importlib.util.module_from_spec(INSTALL_TOOL_SHIMS_SPEC) +INSTALL_TOOL_SHIMS_SPEC.loader.exec_module(install_tool_shims) class CopyPublishedToolsTest(unittest.TestCase): @@ -66,6 +74,66 @@ def test_tool_allowlist_restricts_installed_tools(self) -> None: self.assertTrue((target / "research" / "websearch" / "pyproject.toml").exists()) self.assertFalse((target / "productivity" / "linear").exists()) + def test_tool_allowlist_matches_project_name_when_copying_published_tools(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + published = root / "published" + target = root / "target" + + composio = published / "productivity" / "composio" + composio.mkdir(parents=True) + (composio / "pyproject.toml").write_text( + '[project]\nname = "centaur-composio-tool"\n' + ) + linear = published / "productivity" / "linear" + linear.mkdir(parents=True) + (linear / "pyproject.toml").write_text('[project]\nname = "linear"\n') + + with mock.patch.dict("os.environ", {"TOOL_ALLOWLIST": "centaur-composio-tool"}): + install_tool_shims._copy_published_tools(target, published) + + self.assertTrue((target / "productivity" / "composio" / "pyproject.toml").exists()) + self.assertFalse((target / "productivity" / "linear").exists()) + + def test_tool_blocklist_matches_project_name_when_copying_published_tools(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + published = root / "published" + target = root / "target" + + composio = published / "productivity" / "composio" + composio.mkdir(parents=True) + (composio / "pyproject.toml").write_text( + '[project]\nname = "centaur-composio-tool"\n' + ) + linear = published / "productivity" / "linear" + linear.mkdir(parents=True) + (linear / "pyproject.toml").write_text('[project]\nname = "linear"\n') + + with mock.patch.dict("os.environ", {"TOOL_BLOCKLIST": "centaur-composio-tool"}): + install_tool_shims._copy_published_tools(target, published) + + self.assertFalse((target / "productivity" / "composio").exists()) + self.assertTrue((target / "productivity" / "linear" / "pyproject.toml").exists()) + + def test_malformed_project_table_falls_back_to_directory_matching(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + published = root / "published" + target = root / "target" + + broken = published / "research" / "broken" + broken.mkdir(parents=True) + (broken / "pyproject.toml").write_text('project = "not-a-table"\n') + + stderr = io.StringIO() + with mock.patch.dict("os.environ", {"TOOL_ALLOWLIST": "broken"}): + with contextlib.redirect_stderr(stderr): + install_tool_shims._copy_published_tools(target, published) + + self.assertTrue((target / "research" / "broken" / "pyproject.toml").exists()) + self.assertIn("warning: invalid [project] table", stderr.getvalue()) + def test_unset_allowlist_installs_all_tools(self) -> None: with tempfile.TemporaryDirectory() as tmp: root = Path(tmp)