diff --git a/strix/agents/factory.py b/strix/agents/factory.py index 6a5af17b..65642675 100644 --- a/strix/agents/factory.py +++ b/strix/agents/factory.py @@ -55,7 +55,7 @@ from strix.tools.web_search.tool import web_search if TYPE_CHECKING: - from collections.abc import Awaitable, Callable + from collections.abc import Awaitable, Callable, Sequence from agents import RunContextWrapper from agents.tool import FunctionToolResult @@ -349,6 +349,48 @@ _BASE_TOOLS: tuple[Tool, ...] = ( ) +# Extra tools registered for scan agents. Mirrors +# ``strix.runtime.backends.register_backend``: register before the first +# ``build_strix_agent`` call and every agent (root + children) gets them. +_EXTRA_TOOLS: list[Tool] = [] + + +def _ensure_unique_tool_names(tools: Sequence[Tool]) -> None: + seen: set[str] = set() + duplicates: set[str] = set() + for tool in tools: + if tool.name in seen: + duplicates.add(tool.name) + seen.add(tool.name) + if duplicates: + msg = f"Agent tools must have unique names: {sorted(duplicates)}" + raise ValueError(msg) + + +def register_agent_tools(*tools: Tool) -> None: + """Register tools for every scan agent built afterwards. + + Tools are added to both root and child agents, after the base set and + before the lifecycle tool (``finish_scan`` / ``agent_finish``). Duplicate + tool objects are ignored so repeated imports don't double-register. + """ + new_tools: list[Tool] = [] + for tool in tools: + if tool not in _EXTRA_TOOLS and tool not in new_tools: + new_tools.append(tool) + + _ensure_unique_tool_names([*_BASE_TOOLS, *_EXTRA_TOOLS, *new_tools, finish_scan, agent_finish]) + + for tool in new_tools: + _EXTRA_TOOLS.append(tool) + logger.info("Registered extra agent tool: %s", getattr(tool, "name", tool)) + + +def registered_agent_tools() -> tuple[Tool, ...]: + """Return the currently registered scan-agent tools.""" + return tuple(_EXTRA_TOOLS) + + def build_strix_agent( *, name: str = "strix", @@ -359,26 +401,37 @@ def build_strix_agent( interactive: bool = False, chat_completions_tools: bool = False, system_prompt_context: dict[str, Any] | None = None, + extra_tools: Sequence[Tool] | None = None, + instructions_override: str | None = None, ) -> SandboxAgent[Any]: """Build a SandboxAgent for either root or child use. Args: chat_completions_tools: Wrap SDK custom tools as function tools when the selected backend cannot accept Responses custom tools. + extra_tools: Additional tools for this scan agent only, on top of any + registered via ``register_agent_tools``. + instructions_override: Use this verbatim as the system prompt instead + of rendering the built-in scan prompt. """ - instructions = render_system_prompt( - skills=skills, - scan_mode=scan_mode, - is_whitebox=is_whitebox, - is_root=is_root, - interactive=interactive, - system_prompt_context=system_prompt_context, - ) - - if is_root: - tools: list[Tool] = [*_BASE_TOOLS, finish_scan] + if instructions_override is not None: + instructions = instructions_override else: - tools = [*_BASE_TOOLS, agent_finish] + instructions = render_system_prompt( + skills=skills, + scan_mode=scan_mode, + is_whitebox=is_whitebox, + is_root=is_root, + interactive=interactive, + system_prompt_context=system_prompt_context, + ) + + agent_tools = [*_EXTRA_TOOLS, *(extra_tools or [])] + if is_root: + tools: list[Tool] = [*_BASE_TOOLS, *agent_tools, finish_scan] + else: + tools = [*_BASE_TOOLS, *agent_tools, agent_finish] + _ensure_unique_tool_names(tools) logger.info( "Built %s agent '%s' (skills=%d, tools=%d, scan_mode=%s, whitebox=%s)", diff --git a/tests/test_agent_tool_registration.py b/tests/test_agent_tool_registration.py new file mode 100644 index 00000000..8c7d67df --- /dev/null +++ b/tests/test_agent_tool_registration.py @@ -0,0 +1,88 @@ +"""Tests for scan-agent tool registration in factory.""" + +from __future__ import annotations + +import pytest +from agents.tool import FunctionTool + +from strix.agents import factory + + +def _tool(name: str) -> FunctionTool: + return FunctionTool( + name=name, + description="test tool", + params_json_schema={"type": "object", "properties": {}, "additionalProperties": False}, + on_invoke_tool=lambda _ctx, _inp: "ok", + ) + + +@pytest.fixture(autouse=True) +def _reset_registry() -> object: + saved = list(factory._EXTRA_TOOLS) + factory._EXTRA_TOOLS.clear() + try: + yield + finally: + factory._EXTRA_TOOLS[:] = saved + + +def test_register_agent_tools_is_deduped() -> None: + tool = _tool("dup") + factory.register_agent_tools(tool) + factory.register_agent_tools(tool) + assert factory.registered_agent_tools() == (tool,) + + +def test_registered_tools_appear_before_lifecycle_tool() -> None: + tool = _tool("extra") + factory.register_agent_tools(tool) + + root = factory.build_strix_agent(is_root=True) + child = factory.build_strix_agent(is_root=False) + + root_names = [t.name for t in root.tools] + child_names = [t.name for t in child.tools] + + assert root_names[-2:] == ["extra", "finish_scan"] + assert child_names[-2:] == ["extra", "agent_finish"] + + +def test_per_call_extra_tools_stack_with_registry() -> None: + factory.register_agent_tools(_tool("registered")) + + agent = factory.build_strix_agent(is_root=True, extra_tools=[_tool("per_call")]) + names = [t.name for t in agent.tools] + + assert "registered" in names + assert "per_call" in names + assert names[-1] == "finish_scan" + + +def test_register_agent_tools_rejects_duplicate_names() -> None: + factory.register_agent_tools(_tool("same_name")) + + with pytest.raises(ValueError, match="same_name"): + factory.register_agent_tools(_tool("same_name")) + + +def test_per_call_extra_tools_reject_duplicate_registered_names() -> None: + factory.register_agent_tools(_tool("same_name")) + + with pytest.raises(ValueError, match="same_name"): + factory.build_strix_agent(is_root=True, extra_tools=[_tool("same_name")]) + + +def test_instructions_override_is_used_verbatim() -> None: + custom = "You are a scan agent. Follow the provided scope." + + agent = factory.build_strix_agent(is_root=True, instructions_override=custom) + + assert agent.instructions == custom + + +def test_no_override_renders_builtin_prompt() -> None: + agent = factory.build_strix_agent(is_root=True) + + assert isinstance(agent.instructions, str) + assert agent.instructions != ""