Allow scan agent tool registration

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)",

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 != ""