From 0bb26f4f2a3c8602be005b93255a0ad257fef80e Mon Sep 17 00:00:00 2001 From: Ian Date: Tue, 20 Oct 2026 19:37:38 -0400 Subject: [PATCH] feat(agents): let callers choose the report tool instead of removing tools --- strix/agents/factory.py | 18 +++++++++--------- strix/agents/prompt.py | 23 ++++++++++++++--------- strix/core/runner.py | 9 +++++++++ tests/test_agent_tool_registration.py | 16 +++++++--------- 4 files changed, 39 insertions(+), 27 deletions(-) diff --git a/strix/agents/factory.py b/strix/agents/factory.py index 4c9f81dac..f8a412439 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -670,7 +670,6 @@ _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: @@ -704,11 +703,6 @@ 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) @@ -730,6 +724,7 @@ def build_strix_agent( instructions_override: str | None = None, supports_images: bool = True, finish_tool: Tool = finish_scan, + report_tool: Tool = create_vulnerability_report, ) -> SandboxAgent[Any]: """Build a SandboxAgent for either root or child use. @@ -743,6 +738,8 @@ def build_strix_agent( instructions_override: Use this verbatim as the system prompt instead of rendering the built-in scan prompt. finish_tool: The tool that ends the run, given to the root agent. + report_tool: The tool every agent files findings with, in place of + ``create_vulnerability_report``. """ if instructions_override is not None: instructions = instructions_override @@ -756,17 +753,18 @@ def build_strix_agent( interactive=interactive, system_prompt_context=system_prompt_context, supports_images=supports_images, + report_tool=report_tool.name, ) agent_tools = [*_EXTRA_TOOLS, *(extra_tools or [])] if interactive: # Yielding to the user is only meaningful when one is attached. agent_tools.append(wait_for_user) + base_tools = [report_tool if t is create_vulnerability_report else t for t in _BASE_TOOLS] if is_root: - tools: list[Tool] = [*_BASE_TOOLS, *agent_tools, finish_tool] + 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] + tools = [*base_tools, *agent_tools, agent_finish] _ensure_unique_tool_names(tools) tools = [ _with_bounded_result(_with_strictness(_with_coerced_arguments(tool), strict_tool_schemas)) @@ -819,6 +817,7 @@ def make_child_factory( strict_tool_schemas: bool = True, system_prompt_context: dict[str, Any] | None = None, supports_images: bool = True, + report_tool: Tool = create_vulnerability_report, ) -> Any: """Return the runner-owned builder used by ``spawn_child_agent``. @@ -840,6 +839,7 @@ def make_child_factory( strict_tool_schemas=strict_tool_schemas, system_prompt_context=system_prompt_context, supports_images=supports_images, + report_tool=report_tool, ) return _factory diff --git a/strix/agents/prompt.py b/strix/agents/prompt.py index c45831aae..d6ae1dbc7 100644 --- a/strix/agents/prompt.py +++ b/strix/agents/prompt.py @@ -86,6 +86,7 @@ def render_system_prompt( system_prompt_context: dict[str, Any] | None = None, include_scope: bool = True, supports_images: bool = True, + report_tool: str = "create_vulnerability_report", ) -> str: """Render the system prompt. Returns empty string on template failure. @@ -125,15 +126,19 @@ def render_system_prompt( # Skills every agent of this kind loads come first, so siblings share them # as a cached prefix; the ones the caller asked for vary and go after. - rendered = env.get_template("system_prompt.jinja").render( - shared_skill_names=[name for name in skill_content if name in shared], - requested_skill_names=[name for name in skill_content if name not in shared], - available_skills=get_available_skills(), - interactive=interactive, - is_root=is_root, - system_prompt_context=system_prompt_context or {}, - include_scope=include_scope, - **skill_content, + rendered = ( + env.get_template("system_prompt.jinja") + .render( + shared_skill_names=[name for name in skill_content if name in shared], + requested_skill_names=[name for name in skill_content if name not in shared], + available_skills=get_available_skills(), + interactive=interactive, + is_root=is_root, + system_prompt_context=system_prompt_context or {}, + include_scope=include_scope, + **skill_content, + ) + .replace("create_vulnerability_report", report_tool) ) except Exception: logger.exception("render_system_prompt failed; returning empty prompt") diff --git a/strix/core/runner.py b/strix/core/runner.py index ae6dc08b3..f2482b733 100644 --- a/strix/core/runner.py +++ b/strix/core/runner.py @@ -54,6 +54,7 @@ from strix.tools.output_store import ( WORKSPACE_SPILL_DIR, configure_spill_writer, ) +from strix.tools.reporting.tool import create_vulnerability_report if TYPE_CHECKING: @@ -172,6 +173,7 @@ def _compose_root_instructions_override( interactive: bool, system_prompt_context: dict[str, Any], supports_images: bool, + report_tool: str, ) -> str | None: if root_instructions_override is None: return None @@ -186,6 +188,7 @@ def _compose_root_instructions_override( system_prompt_context=system_prompt_context, include_scope=False, supports_images=supports_images, + report_tool=report_tool, ) return ( f"{base_instructions}\n\n" @@ -218,6 +221,7 @@ async def run_strix_scan( mcp_connection_requests: list[McpConnectionRequest] | None = None, mcp_status_sink: McpStatusSink | None = None, root_finish_tool: Tool = finish_scan, + report_tool: Tool = create_vulnerability_report, ) -> RunResultBase | None: """Run or resume one Strix scan against a sandbox. @@ -241,6 +245,8 @@ async def run_strix_scan( way the engine does the connecting, so the caller passes inert configs plus metadata and never live sessions. ``root_finish_tool`` is the tool the root agent ends the run with. + ``report_tool`` is the tool every agent files findings with, in place of + ``create_vulnerability_report``. """ def report(phase: str) -> None: @@ -495,6 +501,7 @@ async def run_strix_scan( interactive=interactive, system_prompt_context=root_context, supports_images=supports_images, + report_tool=report_tool.name, ) root_agent = build_strix_agent( @@ -511,6 +518,7 @@ async def run_strix_scan( instructions_override=root_instructions, supports_images=supports_images, finish_tool=root_finish_tool, + report_tool=report_tool, ) if not is_resume: @@ -531,6 +539,7 @@ async def run_strix_scan( strict_tool_schemas=strict_tool_schemas, system_prompt_context=scope_context, supports_images=supports_images, + report_tool=report_tool, ) async def spawn_child_agent(**kwargs: Any) -> dict[str, Any]: diff --git a/tests/test_agent_tool_registration.py b/tests/test_agent_tool_registration.py index 822c0b7f6..9af888689 100644 --- a/tests/test_agent_tool_registration.py +++ b/tests/test_agent_tool_registration.py @@ -31,15 +31,11 @@ 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: @@ -97,13 +93,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") - +def test_report_tool_replaces_create_vulnerability_report() -> None: for is_root in (True, False): - names = [t.name for t in factory.build_strix_agent(is_root=is_root).tools] + agent = factory.build_strix_agent(is_root=is_root, report_tool=_tool("file_pr_finding")) + names = [t.name for t in agent.tools] + assert "file_pr_finding" in names assert "create_vulnerability_report" not in names - assert "create_dependency_report" in names + assert isinstance(agent.instructions, str) + assert "file_pr_finding" in agent.instructions + assert "create_vulnerability_report" not in agent.instructions def test_instructions_override_is_used_verbatim() -> None: