diff --git a/src/entrypoints/headless.py b/src/entrypoints/headless.py index fa6eeaa7..0282425b 100644 --- a/src/entrypoints/headless.py +++ b/src/entrypoints/headless.py @@ -173,11 +173,22 @@ def run_headless(options: HeadlessOptions) -> int: session = Session.create(provider_name, getattr(provider, "model", model or "")) tool_registry = build_default_registry(provider=provider) - if options.allowed_tools: - allow = {name.lower() for name in options.allowed_tools} + # Canonicalize BOTH sets up front (before either filter runs) so an alias + # form (e.g. --disallowed-tools KillShell) resolves while its tool is still + # registered. + allow = ( + tool_registry.canonicalize_names(options.allowed_tools) + if options.allowed_tools + else None + ) + deny = ( + tool_registry.canonicalize_names(options.disallowed_tools) + if options.disallowed_tools + else None + ) + if allow is not None: _filter_registry(tool_registry, keep=lambda n: n.lower() in allow) - if options.disallowed_tools: - deny = {name.lower() for name in options.disallowed_tools} + if deny is not None: _filter_registry(tool_registry, keep=lambda n: n.lower() not in deny) # (workspace_root already resolved above, before the prefetch kick.) @@ -841,7 +852,13 @@ def _restore() -> None: def _filter_registry(registry, *, keep) -> None: - """In-place best-effort filter of a ToolRegistry.""" + """In-place best-effort filter of a ToolRegistry. + + Drops every tool for which ``keep(name)`` is False so that + ``--allowed-tools`` / ``--disallowed-tools`` remove the tool from the pool + the model sees (schemas are emitted from ``registry.list_tools()``), not + just block it at execution time. + """ try: entries = list(registry.list_tools()) @@ -851,10 +868,10 @@ def _filter_registry(registry, *, keep) -> None: name = getattr(tool, "name", "") if not keep(name): try: - registry.unregister(name) + registry.remove_tool(name) except Exception: - # Registry may not support unregistration; fall back to - # marking the tool disallowed through ToolContext. + # Best-effort: a registry that cannot drop the tool leaves it + # in the pool rather than aborting the whole filter. continue diff --git a/src/server/agent_server.py b/src/server/agent_server.py index f20ccc84..b0683599 100644 --- a/src/server/agent_server.py +++ b/src/server/agent_server.py @@ -1089,11 +1089,21 @@ def _do_set_provider(self, request_id: object, name: object) -> None: provider = provider_cls(api_key=api_key, base_url=provider_cfg.get("base_url"), model=model) registry = build_default_registry(provider=provider) cfg = self.config - if cfg.allowed_tools: - allow = {n.lower() for n in cfg.allowed_tools} + # Canonicalize up front so alias-form flags (e.g. KillShell -> + # TaskStop) resolve while the tool is still registered. + allow = ( + registry.canonicalize_names(cfg.allowed_tools) + if cfg.allowed_tools + else None + ) + deny = ( + registry.canonicalize_names(cfg.disallowed_tools) + if cfg.disallowed_tools + else None + ) + if allow is not None: _filter_registry(registry, keep=lambda n: n.lower() in allow) - if cfg.disallowed_tools: - deny = {n.lower() for n in cfg.disallowed_tools} + if deny is not None: _filter_registry(registry, keep=lambda n: n.lower() not in deny) if self._mcp_runtime is not None: # keep MCP tools across the switch for mtool in self._mcp_runtime.tools: @@ -4202,6 +4212,12 @@ def _with_ultracode_reminder(prompt): def _filter_registry(registry, *, keep) -> None: + """Drop every tool for which ``keep(name)`` is False. + + Backs ``--allowed-tools`` / ``--disallowed-tools``: removing the tool from + the registry keeps its schema out of the ``tools=`` param sent to the + model, not just blocked at execution time. + """ try: entries = list(registry.list_tools()) except Exception: # noqa: BLE001 @@ -4210,7 +4226,7 @@ def _filter_registry(registry, *, keep) -> None: name = getattr(tool, "name", "") if not keep(name): try: - registry.unregister(name) + registry.remove_tool(name) except Exception: # noqa: BLE001 continue diff --git a/src/server/mcp_runtime.py b/src/server/mcp_runtime.py index 3e936ba5..8972d7aa 100644 --- a/src/server/mcp_runtime.py +++ b/src/server/mcp_runtime.py @@ -298,7 +298,7 @@ def _apply_refreshed_tools( server's slice of self.tools/self.servers, and return ``(removed_full_names, new_wrapped_tools)`` so the caller can swap them in the live agent registry. Returns the FULL - ``mcp__server__tool`` names removed (what registry.unregister keys on). + ``mcp__server__tool`` names removed (what registry.remove_tool keys on). """ from src.services.mcp.mcp_string_utils import build_mcp_tool_name diff --git a/src/tool_system/registry.py b/src/tool_system/registry.py index e06cc6d3..10a9f3d6 100644 --- a/src/tool_system/registry.py +++ b/src/tool_system/registry.py @@ -137,29 +137,53 @@ def remove_tool(self, name: str) -> bool: when a server sends notifications/tools/list_changed — the registry was otherwise append-only, so a re-fetch couldn't reach the agent. - NB: deliberately NOT named ``unregister`` — ``_filter_registry`` - (--allowedTools/--disallowedTools, agent_server.py + headless.py) - already calls a non-existent ``registry.unregister`` inside a - try/except as a silent no-op, so its registry-level filtering has - never actually run (the live enforcement is the deny-rule path at - assembly). Adding an ``unregister`` here would silently ACTIVATE that - untested, security-relevant path; a distinct name keeps this MCP - change scoped. Activating that filtering is a separate follow-up.""" - key = name.lower() - tool = self._by_name.pop(key, None) + Also backs ``_filter_registry`` (--allowed-tools/--disallowed-tools in + agent_server.py + headless.py): those call sites now call this method + directly. They historically called a non-existent + ``registry.unregister`` inside a try/except, so the registry-level + filtering silently no-op'd — the flags removed nothing from the pool + the model saw. Both paths only ever REMOVE tools, so activating them + can only narrow the toolset. + + Resolving by ``name`` accepts a primary name OR an alias, then drops + the tool from ``_tools`` and EVERY ``_by_name`` key that still points + at it (canonical + aliases). Removing by an alias therefore also + clears the canonical key, so the tool can't stay dispatch-reachable + through a sibling key.""" + tool = self._by_name.get(name.lower()) if tool is None: return False self._tools = [t for t in self._tools if t is not tool] - for alias in getattr(tool, "aliases", ()) or (): - # Only drop the alias if it still points at THIS tool (another - # tool may have claimed it — don't clobber that). - if self._by_name.get(alias.lower()) is tool: - self._by_name.pop(alias.lower(), None) + for key in (tool.name, *(getattr(tool, "aliases", ()) or ())): + # Only drop a key that still points at THIS tool (another tool may + # have claimed an alias — don't clobber that). + if self._by_name.get(key.lower()) is tool: + self._by_name.pop(key.lower(), None) return True def get(self, name: str) -> Tool | None: return self._by_name.get(name.lower()) + def canonicalize_names(self, names: Iterable[str]) -> set[str]: + """Resolve each requested name to its tool's lowercased PRIMARY name. + + Backs ``--allowed-tools`` / ``--disallowed-tools`` so they match whether + the caller passes a tool's primary name or one of its aliases + (e.g. ``KillShell`` -> ``TaskStop``, ``Task`` -> ``Agent``). + ``list_tools()`` only yields primary names and ``remove_tool`` cleans a + tool's aliases with it, so canonicalizing the request set to primary + names is what makes alias-form flags actually filter. Unknown names + pass through lowercased — they simply match nothing.""" + out: set[str] = set() + for name in names: + if not name or not name.strip(): + # Skip blanks so a stray "" can't become a match-nothing + # allowlist that removes every tool. + continue + tool = self.get(name) + out.add((tool.name if tool else name).lower()) + return out + def list_tools(self) -> Tools: """Tools visible to the agent — excludes disabled-MCP-server tools.""" if not self.disabled_servers: diff --git a/tests/test_headless_tool_filter.py b/tests/test_headless_tool_filter.py new file mode 100644 index 00000000..fa45c3f6 --- /dev/null +++ b/tests/test_headless_tool_filter.py @@ -0,0 +1,144 @@ +"""Regression tests for ``--allowedTools`` / ``--disallowedTools`` filtering. + +Both the headless (``--print``) and agent-server paths expose a +``_filter_registry`` helper that backs the ``--allowedTools`` / +``--disallowedTools`` CLI flags. It historically called a non-existent +``registry.unregister`` inside a bare ``try/except``, so the filtering was a +silent no-op: the flags removed nothing from the pool the model saw (tool +schemas are emitted from ``registry.list_tools()``). These tests lock in that +the helper now calls the real ``ToolRegistry.remove_tool`` so the flags take +effect. +""" + +from __future__ import annotations + +from src.entrypoints import headless as headless_mod +from src.server import agent_server as server_mod +from src.tool_system.defaults import build_default_registry + + +def _names(registry) -> set[str]: + return {t.name for t in registry.list_tools()} + + +def test_registry_lacks_unregister_but_has_remove_tool(): + """The bug was calling ``unregister`` (absent) instead of ``remove_tool``. + + Guards against a future rename silently reintroducing the dead path. + """ + registry = build_default_registry(provider="anthropic") + assert not hasattr(registry, "unregister") + assert hasattr(registry, "remove_tool") + + +def test_headless_filter_denylist_removes_from_pool(): + registry = build_default_registry(provider="anthropic") + assert "AskUserQuestion" in _names(registry) + + deny = {"askuserquestion"} + headless_mod._filter_registry( + registry, keep=lambda n: n.lower() not in deny + ) + + remaining = _names(registry) + assert "AskUserQuestion" not in remaining, ( + "disallowed tool must be dropped from list_tools() (the schema source), " + "not left in the pool the model sees" + ) + # Untargeted tools survive. + assert "Bash" in remaining + assert "Read" in remaining + + +def test_headless_filter_allowlist_keeps_only_allowed(): + registry = build_default_registry(provider="anthropic") + keep_set = {"bash", "read", "write", "edit"} + headless_mod._filter_registry(registry, keep=lambda n: n.lower() in keep_set) + + remaining = {n.lower() for n in _names(registry)} + assert remaining == keep_set, remaining + + +def test_agent_server_filter_matches_headless_behavior(): + """The agent-server (TUI/interactive) path shares the same contract.""" + registry = build_default_registry(provider="anthropic") + deny = {"workflow", "croncreate"} + server_mod._filter_registry(registry, keep=lambda n: n.lower() not in deny) + + remaining = {n.lower() for n in _names(registry)} + assert "workflow" not in remaining + assert "croncreate" not in remaining + assert "bash" in remaining + + +def test_filter_is_idempotent_and_survives_unknown_names(): + registry = build_default_registry(provider="anthropic") + before = _names(registry) + # Denying a name that isn't registered must not raise or change anything. + headless_mod._filter_registry(registry, keep=lambda n: n.lower() != "not_a_tool") + assert _names(registry) == before + + +def test_denylist_by_alias_removes_tool_from_schema_and_dispatch(): + """--disallowedTools accepts an ALIAS and still removes the real tool. + + ``TaskStop`` carries the ``KillShell`` alias. Denying the alias must drop + ``TaskStop`` from ``list_tools()`` (the schema the model sees) AND make it + unreachable via ``get()`` (dispatch), not leave it half-present. + """ + registry = build_default_registry(provider="anthropic") + assert registry.get("KillShell") is not None # alias resolves pre-filter + assert "TaskStop" in _names(registry) + + deny = registry.canonicalize_names(["KillShell"]) + assert deny == {"taskstop"} # alias resolved to primary + headless_mod._filter_registry(registry, keep=lambda n: n.lower() not in deny) + + assert "TaskStop" not in _names(registry), "alias-denied tool still in schema" + assert registry.get("TaskStop") is None, "primary still dispatch-reachable" + assert registry.get("KillShell") is None, "alias still dispatch-reachable" + + +def test_allowlist_by_alias_keeps_the_real_tool(): + """--allowedTools with an alias keeps its tool (and drops the rest).""" + registry = build_default_registry(provider="anthropic") + allow = registry.canonicalize_names(["Task", "Bash"]) # Task -> Agent + assert allow == {"agent", "bash"} + headless_mod._filter_registry(registry, keep=lambda n: n.lower() in allow) + + remaining = {n.lower() for n in _names(registry)} + assert remaining == {"agent", "bash"} + + +def test_canonicalize_passes_unknown_names_through(): + registry = build_default_registry(provider="anthropic") + assert registry.canonicalize_names(["Bash", "NotATool"]) == {"bash", "notatool"} + + +def test_remove_tool_by_alias_clears_canonical_dispatch_key(): + """remove_tool given an ALIAS must also drop the canonical key. + + Defense-in-depth for a now load-bearing method: a future caller passing an + alias must not leave the tool dispatch-reachable via its primary name. + """ + registry = build_default_registry(provider="anthropic") + assert registry.get("TaskStop") is not None + assert registry.remove_tool("KillShell") is True # removed by ALIAS + assert "TaskStop" not in _names(registry) + assert registry.get("TaskStop") is None # canonical key cleared + assert registry.get("KillShell") is None # alias key cleared + + +def test_canonicalize_skips_blanks_so_allowlist_cannot_wipe_all(): + """A stray "" must not become a match-nothing allowlist removing everything.""" + registry = build_default_registry(provider="anthropic") + assert registry.canonicalize_names(["", " ", "Bash"]) == {"bash"} + + allow = registry.canonicalize_names([""]) + assert allow == set() + # An empty allow set means "no allowlist constraint given" at the call + # sites (they gate on truthiness), so nothing is filtered here. + before = _names(registry) + if allow: + headless_mod._filter_registry(registry, keep=lambda n: n.lower() in allow) + assert _names(registry) == before