strix/tests/test_agent_tool_registration.py

"""Tests for scan-agent tool registration in factory."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

import pytest
from agents.tool import FunctionTool

from strix.agents import factory


if TYPE_CHECKING:
    from agents.tool_context import ToolContext


def _tool(name: str) -> FunctionTool:
    # A per-tool closure keeps two same-named tools unequal, which is what the
    # duplicate-name tests exercise.
    async def invoke(_ctx: ToolContext[Any], _input: str) -> str:
        return "ok"

    return FunctionTool(
        name=name,
        description="test tool",
        params_json_schema={"type": "object", "properties": {}, "additionalProperties": False},
        on_invoke_tool=invoke,
    )


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


def test_wait_for_user_is_interactive_only() -> None:
    """Yielding to the user is meaningless when no user is attached."""
    interactive = factory.build_strix_agent(is_root=True, interactive=True)
    autonomous = factory.build_strix_agent(is_root=True, interactive=False)

    assert "wait_for_user" in [t.name for t in interactive.tools]
    assert "wait_for_user" not in [t.name for t in autonomous.tools]


def test_wait_for_agents_is_available_in_both_modes() -> None:
    for interactive in (True, False):
        agent = factory.build_strix_agent(is_root=True, interactive=interactive)
        assert "wait_for_agents" in [t.name for t in agent.tools]


def test_strict_tool_schemas_can_be_disabled_per_route() -> None:
    """Claude routes cap strict tools; the toolset must be sendable without strict."""
    agent = factory.build_strix_agent(is_root=True, strict_tool_schemas=False)

    function_tools = [t for t in agent.tools if isinstance(t, FunctionTool)]
    assert function_tools
    assert not any(t.strict_json_schema for t in function_tools)


def test_disabling_strict_leaves_shared_tools_untouched() -> None:
    factory.build_strix_agent(is_root=True, strict_tool_schemas=False)
    agent = factory.build_strix_agent(is_root=True)

    assert any(t.strict_json_schema for t in agent.tools if isinstance(t, FunctionTool))


def _untyped_properties(schema: dict[str, Any], path: str) -> list[str]:
    untyped = []
    for key, spec in (schema.get("properties") or {}).items():
        if "type" not in spec:
            untyped.append(f"{path}.{key}")
        untyped += _untyped_properties(spec, f"{path}.{key}")
        if isinstance(spec.get("items"), dict):
            untyped += _untyped_properties(spec["items"], f"{path}.{key}[]")
    return untyped


@pytest.mark.parametrize("strict", [True, False])
def test_chat_completions_tool_parameters_all_declare_a_type(strict: bool) -> None:
    """Ollama's tool templates index ``$prop.Type`` and fail on a bare ``anyOf``."""
    agent = factory.build_strix_agent(
        is_root=True, chat_completions_tools=True, strict_tool_schemas=strict, interactive=True
    )

    untyped = [
        path
        for t in agent.tools
        if isinstance(t, FunctionTool)
        for path in _untyped_properties(t.params_json_schema, t.name)
    ]
    assert untyped == []


def test_nullable_parameter_keeps_its_variant_schema() -> None:
    spec = {
        "anyOf": [{"type": "array", "items": {"type": "string"}}, {"type": "null"}],
        "default": None,
        "description": "Filter by tag.",
    }

    assert factory._typed_schema({"type": "object", "properties": {"tags": spec}}) == {
        "type": "object",
        "properties": {
            "tags": {
                "type": ["array", "null"],
                "items": {"type": "string"},
                "default": None,
                "description": "Filter by tag.",
            }
        },
    }


@pytest.mark.parametrize(
    "spec",
    [
        {"anyOf": [{"type": "string", "enum": ["a", "b"]}, {"type": "null"}]},
        {"anyOf": [{"type": "string", "const": "a"}, {"type": "null"}]},
        {"anyOf": [{"type": "string", "format": "uri"}, {"type": "array", "format": "x"}]},
        {"anyOf": [{"$ref": "#/$defs/Mode"}, {"type": "null"}]},
    ],
)
def test_unmergeable_unions_are_left_untouched(spec: dict[str, Any]) -> None:
    schema = {"type": "object", "properties": {"value": spec}}

    assert factory._typed_schema(schema) == schema


def test_typing_parameters_leaves_shared_tools_untouched() -> None:
    factory.build_strix_agent(is_root=True, chat_completions_tools=True)
    agent = factory.build_strix_agent(is_root=True)
    list_notes = next(t for t in agent.tools if t.name == "list_notes")

    assert isinstance(list_notes, FunctionTool)
    assert "anyOf" in list_notes.params_json_schema["properties"]["tags"]


@pytest.mark.parametrize(
    ("tool_name", "output", "completed"),
    [
        ("finish_pr_review", '{"success": true, "review_completed": true}', True),
        ("finish_pr_review", '{"success": false, "error": "agents_still_active"}', False),
        ("finish_scan", '{"success": true, "scan_completed": true}', True),
        ("think", '{"success": true, "review_completed": true}', False),
    ],
)
def test_lifecycle_tool_completed(tool_name: str, output: str, completed: bool) -> None:
    assert factory._lifecycle_tool_completed(tool_name, output) is completed