feat(agents): let callers remove base tools from scan agents

This commit is contained in:
Ian 2026-10-20 19:16:45 -04:00
parent aa2036f229
commit 168826a335
2 changed files with 20 additions and 0 deletions

View file

@ -670,6 +670,7 @@ _BASE_TOOLS: tuple[Tool, ...] = (
# ``strix.runtime.backends.register_backend``: register before the first
# ``build_strix_agent`` call and every agent (root + children) gets them.
_EXTRA_TOOLS: list[Tool] = []
_REMOVED_TOOL_NAMES: set[str] = set()
def _ensure_unique_tool_names(tools: Sequence[Tool]) -> None:
@ -703,6 +704,11 @@ def register_agent_tools(*tools: Tool) -> None:
logger.info("Registered extra agent tool: %s", getattr(tool, "name", tool))
def remove_agent_tools(*names: str) -> None:
"""Leave these base tools out of every scan agent built afterwards."""
_REMOVED_TOOL_NAMES.update(names)
def registered_agent_tools() -> tuple[Tool, ...]:
"""Return the currently registered scan-agent tools."""
return tuple(_EXTRA_TOOLS)
@ -760,6 +766,7 @@ def build_strix_agent(
tools: list[Tool] = [*_BASE_TOOLS, *agent_tools, finish_tool]
else:
tools = [*_BASE_TOOLS, *agent_tools, agent_finish]
tools = [tool for tool in tools if tool.name not in _REMOVED_TOOL_NAMES]
_ensure_unique_tool_names(tools)
tools = [
_with_bounded_result(_with_strictness(_with_coerced_arguments(tool), strict_tool_schemas))

View file

@ -31,11 +31,15 @@ def _tool(name: str) -> FunctionTool:
@pytest.fixture(autouse=True)
def _reset_registry() -> object:
saved = list(factory._EXTRA_TOOLS)
saved_removed = set(factory._REMOVED_TOOL_NAMES)
factory._EXTRA_TOOLS.clear()
factory._REMOVED_TOOL_NAMES.clear()
try:
yield
finally:
factory._EXTRA_TOOLS[:] = saved
factory._REMOVED_TOOL_NAMES.clear()
factory._REMOVED_TOOL_NAMES.update(saved_removed)
def test_register_agent_tools_is_deduped() -> None:
@ -93,6 +97,15 @@ def test_per_call_extra_tools_reject_duplicate_registered_names() -> None:
factory.build_strix_agent(is_root=True, extra_tools=[_tool("same_name")])
def test_removed_tools_are_left_out_of_root_and_child() -> None:
factory.remove_agent_tools("create_vulnerability_report")
for is_root in (True, False):
names = [t.name for t in factory.build_strix_agent(is_root=is_root).tools]
assert "create_vulnerability_report" not in names
assert "create_dependency_report" in names
def test_instructions_override_is_used_verbatim() -> None:
custom = "You are a scan agent. Follow the provided scope."