test_tools.py

"""Unit tests for define_tool"""

import json

import pytest
from pydantic import BaseModel, ConfigDict, Field, field_validator

from copilot import define_tool
from copilot.generated.rpc import ExternalToolTextResultForLlm
from copilot.tools import (
    ToolBinaryResult,
    ToolInvocation,
    ToolResult,
    _normalize_result,
    convert_mcp_call_tool_result,
    tool_result_to_external_tool_text_result_for_llm,
)


class TestDefineTool:
    def test_creates_tool_with_correct_name_and_description(self):
        class Params(BaseModel):
            query: str

        @define_tool("search", description="Search for something")
        def search(params: Params, invocation: ToolInvocation) -> str:
            return "result"

        assert search.name == "search"
        assert search.description == "Search for something"
        assert search.handler is not None
        assert search.parameters is not None

    def test_infers_name_from_function(self):
        class Params(BaseModel):
            query: str

        @define_tool(description="Search for something")
        def my_search_tool(params: Params) -> str:
            return "result"

        assert my_search_tool.name == "my_search_tool"

    def test_generates_schema_from_pydantic_model(self):
        class Params(BaseModel):
            city: str = Field(description="City name")
            unit: str = Field(description="Temperature unit")

        @define_tool("get_weather", description="Get weather")
        def get_weather(params: Params, invocation: ToolInvocation) -> str:
            return "sunny"

        schema = get_weather.parameters
        assert schema is not None
        assert schema["type"] == "object"
        assert "city" in schema["properties"]
        assert "unit" in schema["properties"]
        assert schema["properties"]["city"]["description"] == "City name"

    async def test_handler_receives_typed_arguments(self):
        class Params(BaseModel):
            name: str
            count: int

        received_params = None

        @define_tool("test", description="Test tool")
        def test_tool(params: Params, invocation: ToolInvocation) -> str:
            nonlocal received_params
            received_params = params
            return "ok"

        invocation = ToolInvocation(
            session_id="session-1",
            tool_call_id="call-1",
            tool_name="test",
            arguments={"name": "Alice", "count": 42},
        )

        await test_tool.handler(invocation)

        assert received_params is not None
        assert received_params.name == "Alice"
        assert received_params.count == 42

    async def test_handler_receives_invocation(self):
        class Params(BaseModel):
            pass

        received_inv = None

        @define_tool("test", description="Test tool")
        def test_tool(params: Params, invocation: ToolInvocation) -> str:
            nonlocal received_inv
            received_inv = invocation
            return "ok"

        invocation = ToolInvocation(
            session_id="session-123",
            tool_call_id="call-456",
            tool_name="test",
            arguments={},
        )

        await test_tool.handler(invocation)

        assert received_inv.session_id == "session-123"
        assert received_inv.tool_call_id == "call-456"

    async def test_zero_param_handler(self):
        """Handler with no parameters: def handler() -> str"""
        called = False

        @define_tool("test", description="Test tool")
        def test_tool() -> str:
            nonlocal called
            called = True
            return "ok"

        invocation = ToolInvocation(
            session_id="s1",
            tool_call_id="c1",
            tool_name="test",
            arguments={},
        )

        result = await test_tool.handler(invocation)

        assert called
        assert result.text_result_for_llm == "ok"

    async def test_invocation_only_handler(self):
        """Handler with only invocation: def handler(invocation) -> str"""
        received_inv = None

        @define_tool("test", description="Test tool")
        def test_tool(invocation: ToolInvocation) -> str:
            nonlocal received_inv
            received_inv = invocation
            return "ok"

        invocation = ToolInvocation(
            session_id="s1",
            tool_call_id="c1",
            tool_name="test",
            arguments={},
        )

        await test_tool.handler(invocation)

        assert received_inv is not None
        assert received_inv.session_id == "s1"

    async def test_params_only_handler(self):
        """Handler with only params: def handler(params) -> str"""

        class Params(BaseModel):
            value: str

        received_params = None

        @define_tool("test", description="Test tool")
        def test_tool(params: Params) -> str:
            nonlocal received_params
            received_params = params
            return "ok"

        invocation = ToolInvocation(
            session_id="s1",
            tool_call_id="c1",
            tool_name="test",
            arguments={"value": "hello"},
        )

        await test_tool.handler(invocation)

        assert received_params is not None
        assert received_params.value == "hello"

    async def test_handler_error_is_hidden_from_llm(self):
        class Params(BaseModel):
            pass

        @define_tool("failing", description="A failing tool")
        def failing_tool(params: Params, invocation: ToolInvocation) -> str:
            raise ValueError("secret error message")

        invocation = ToolInvocation(
            session_id="s1",
            tool_call_id="c1",
            tool_name="failing",
            arguments={},
        )

        result = await failing_tool.handler(invocation)

        assert result.result_type == "failure"
        assert "secret error message" not in result.text_result_for_llm
        assert "error" in result.text_result_for_llm.lower()
        # But the actual error is stored internally
        assert result.error == "secret error message"

    async def test_validation_error_is_surfaced_to_llm(self):
        class Params(BaseModel):
            username: str

            @field_validator("username")
            @classmethod
            def check_username(cls, v: str) -> str:
                if v == "admin":
                    raise ValueError("username 'admin' is reserved")
                return v

        @define_tool("validate", description="A validating tool")
        def validating_tool(params: Params) -> str:
            return "ok"

        invocation = ToolInvocation(
            session_id="s1",
            tool_call_id="c1",
            tool_name="validate",
            arguments={"username": "admin"},
        )

        result = await validating_tool.handler(invocation)

        assert result.result_type == "failure"
        assert result.text_result_for_llm.startswith("Invalid tool arguments:")
        assert "username 'admin' is reserved" in result.text_result_for_llm
        # Full detail is retained in the debug field.
        assert result.error is not None

    async def test_validation_error_extra_forbid_includes_field_name(self):
        class Params(BaseModel):
            model_config = ConfigDict(extra="forbid")

            request: str

        @define_tool("strict", description="A strict tool")
        def strict_tool(params: Params) -> str:
            return "ok"

        invocation = ToolInvocation(
            session_id="s1",
            tool_call_id="c1",
            tool_name="strict",
            arguments={"request": "ok", "extra_field": "unexpected"},
        )

        result = await strict_tool.handler(invocation)

        assert result.result_type == "failure"
        assert result.text_result_for_llm.startswith("Invalid tool arguments:")
        # The offending key name is carried in `loc` even though the generic
        # message is "Extra inputs are not permitted".
        assert "extra_field" in result.text_result_for_llm
        assert result.error is not None

    async def test_validation_error_from_handler_body_is_redacted(self):
        class Params(BaseModel):
            pass

        class Internal(BaseModel):
            count: int

        @define_tool("body", description="A tool that validates internally")
        def body_tool(params: Params) -> str:
            Internal.model_validate({"count": "secret-not-an-int"})
            return "ok"

        invocation = ToolInvocation(
            session_id="s1",
            tool_call_id="c1",
            tool_name="body",
            arguments={},
        )

        result = await body_tool.handler(invocation)

        assert result.result_type == "failure"
        # A ValidationError from the handler body must not be surfaced as an
        # argument-validation error; it stays redacted like any other exception.
        assert not result.text_result_for_llm.startswith("Invalid tool arguments:")
        assert "secret-not-an-int" not in result.text_result_for_llm
        assert "error" in result.text_result_for_llm.lower()
        assert result.error is not None

    async def test_function_style_api(self):
        class Params(BaseModel):
            value: str

        tool = define_tool(
            "my_tool",
            description="My tool",
            handler=lambda params, inv: params.value.upper(),
            params_type=Params,
        )

        assert tool.name == "my_tool"
        assert tool.description == "My tool"

        result = await tool.handler(
            ToolInvocation(
                session_id="s",
                tool_call_id="c",
                tool_name="my_tool",
                arguments={"value": "hello"},
            )
        )
        assert result.text_result_for_llm == "HELLO"

    def test_function_style_can_create_declaration_only_tool(self):
        class Params(BaseModel):
            value: str = Field(description="Value to look up")

        tool = define_tool(
            "my_tool",
            description="My tool",
            params_type=Params,
        )

        assert tool.name == "my_tool"
        assert tool.description == "My tool"
        assert tool.handler is None
        assert tool.parameters is not None
        assert tool.parameters["properties"]["value"]["description"] == "Value to look up"

    def test_function_style_requires_name(self):
        class Params(BaseModel):
            value: str

        with pytest.raises(ValueError, match="name is required"):
            define_tool(
                description="My tool",
                handler=lambda params, inv: params.value.upper(),
                params_type=Params,
            )


class TestNormalizeResult:
    def test_none_returns_empty_success(self):
        result = _normalize_result(None)
        assert result.text_result_for_llm == ""
        assert result.result_type == "success"

    def test_string_passes_through(self):
        result = _normalize_result("hello world")
        assert result.text_result_for_llm == "hello world"
        assert result.result_type == "success"

    def test_tool_result_passes_through(self):
        input_result = ToolResult(
            text_result_for_llm="custom",
            result_type="failure",
            error="some error",
        )
        result = _normalize_result(input_result)
        assert result.text_result_for_llm == "custom"
        assert result.result_type == "failure"

    def test_dict_is_json_serialized(self):
        result = _normalize_result({"key": "value", "num": 42})
        parsed = json.loads(result.text_result_for_llm)
        assert parsed == {"key": "value", "num": 42}
        assert result.result_type == "success"

    def test_list_is_json_serialized(self):
        result = _normalize_result(["a", "b", "c"])
        assert result.text_result_for_llm == '["a", "b", "c"]'
        assert result.result_type == "success"

    def test_pydantic_model_is_serialized(self):
        class Response(BaseModel):
            status: str
            count: int

        result = _normalize_result(Response(status="ok", count=5))
        parsed = json.loads(result.text_result_for_llm)
        assert parsed == {"status": "ok", "count": 5}

    def test_list_of_pydantic_models_is_serialized(self):
        class Item(BaseModel):
            name: str
            value: int

        items = [Item(name="a", value=1), Item(name="b", value=2)]
        result = _normalize_result(items)
        parsed = json.loads(result.text_result_for_llm)
        assert parsed == [{"name": "a", "value": 1}, {"name": "b", "value": 2}]
        assert result.result_type == "success"

    def test_raises_for_unserializable_value(self):
        # Functions cannot be JSON serialized
        with pytest.raises(TypeError, match="Failed to serialize"):
            _normalize_result(lambda x: x)


class TestConvertMcpCallToolResult:
    def test_text_only_call_tool_result(self):
        result = convert_mcp_call_tool_result(
            {
                "content": [{"type": "text", "text": "hello"}],
            }
        )
        assert result.text_result_for_llm == "hello"
        assert result.result_type == "success"

    def test_multiple_text_blocks(self):
        result = convert_mcp_call_tool_result(
            {
                "content": [
                    {"type": "text", "text": "line 1"},
                    {"type": "text", "text": "line 2"},
                ],
            }
        )
        assert result.text_result_for_llm == "line 1\nline 2"

    def test_is_error_maps_to_failure(self):
        result = convert_mcp_call_tool_result(
            {
                "content": [{"type": "text", "text": "oops"}],
                "isError": True,
            }
        )
        assert result.result_type == "failure"

    def test_is_error_false_maps_to_success(self):
        result = convert_mcp_call_tool_result(
            {
                "content": [{"type": "text", "text": "ok"}],
                "isError": False,
            }
        )
        assert result.result_type == "success"

    def test_image_content_to_binary(self):
        result = convert_mcp_call_tool_result(
            {
                "content": [{"type": "image", "data": "base64data", "mimeType": "image/png"}],
            }
        )
        assert result.binary_results_for_llm is not None
        assert len(result.binary_results_for_llm) == 1
        assert result.binary_results_for_llm[0].data == "base64data"
        assert result.binary_results_for_llm[0].mime_type == "image/png"
        assert result.binary_results_for_llm[0].type == "image"

    def test_resource_text_to_text_result(self):
        result = convert_mcp_call_tool_result(
            {
                "content": [
                    {
                        "type": "resource",
                        "resource": {"uri": "file:///data.txt", "text": "file contents"},
                    },
                ],
            }
        )
        assert result.text_result_for_llm == "file contents"

    def test_resource_blob_to_binary(self):
        result = convert_mcp_call_tool_result(
            {
                "content": [
                    {
                        "type": "resource",
                        "resource": {
                            "uri": "file:///img.png",
                            "blob": "blobdata",
                            "mimeType": "image/png",
                        },
                    },
                ],
            }
        )
        assert result.binary_results_for_llm is not None
        assert len(result.binary_results_for_llm) == 1
        assert result.binary_results_for_llm[0].data == "blobdata"
        assert result.binary_results_for_llm[0].mime_type == "image/png"
        assert result.binary_results_for_llm[0].description == "file:///img.png"

    def test_resource_blob_defaults_missing_or_empty_mime_type(self):
        result = convert_mcp_call_tool_result(
            {
                "content": [
                    {
                        "type": "resource",
                        "resource": {"uri": "file:///data.bin", "blob": "binarydata"},
                    },
                    {
                        "type": "resource",
                        "resource": {
                            "uri": "file:///empty-mime.bin",
                            "blob": "binarydata2",
                            "mimeType": "",
                        },
                    },
                ],
            }
        )

        assert result.binary_results_for_llm is not None
        assert len(result.binary_results_for_llm) == 2
        assert result.binary_results_for_llm[0].mime_type == "application/octet-stream"
        assert result.binary_results_for_llm[1].mime_type == "application/octet-stream"

    def test_empty_content_array(self):
        result = convert_mcp_call_tool_result({"content": []})
        assert result.text_result_for_llm == ""
        assert result.result_type == "success"

    def test_call_tool_result_dict_is_json_serialized_by_normalize(self):
        """_normalize_result does NOT auto-detect MCP results; it JSON-serializes them."""
        result = _normalize_result({"content": [{"type": "text", "text": "hello"}]})
        parsed = json.loads(result.text_result_for_llm)
        assert parsed == {"content": [{"type": "text", "text": "hello"}]}


class TestToolReferences:
    def test_tool_references_pass_through_normalize(self):
        input_result = ToolResult(
            text_result_for_llm="found 2 tools",
            result_type="success",
            tool_references=["get_weather", "check_status"],
        )
        result = _normalize_result(input_result)
        assert result.tool_references == ["get_weather", "check_status"]

    def test_tool_references_serialized_to_wire(self):
        wire = ExternalToolTextResultForLlm(
            text_result_for_llm="found 2 tools",
            result_type="success",
            tool_references=["get_weather", "check_status"],
        )
        data = wire.to_dict()
        assert data["toolReferences"] == ["get_weather", "check_status"]

    def test_tool_references_omitted_when_none(self):
        wire = ExternalToolTextResultForLlm(
            text_result_for_llm="ok",
            result_type="success",
        )
        assert "toolReferences" not in wire.to_dict()

    def test_tool_references_round_trip_from_wire(self):
        wire = ExternalToolTextResultForLlm.from_dict(
            {
                "textResultForLlm": "found tools",
                "resultType": "success",
                "toolReferences": ["alpha", "beta"],
            }
        )
        assert wire.tool_references == ["alpha", "beta"]


class TestToolResultToExternalToolTextResultForLlm:
    def test_forwards_binary_results_and_session_log(self):
        tool_result = ToolResult(
            text_result_for_llm="screenshot captured",
            binary_results_for_llm=[
                ToolBinaryResult(
                    data="base64data",
                    mime_type="image/png",
                    type="image",
                    description="screenshot.png",
                )
            ],
            session_log="tool execution details",
            tool_telemetry={"duration_ms": 42},
        )

        rpc_result = tool_result_to_external_tool_text_result_for_llm(tool_result)

        assert rpc_result.text_result_for_llm == "screenshot captured"
        assert rpc_result.session_log == "tool execution details"
        assert rpc_result.tool_telemetry == {"duration_ms": 42}
        assert rpc_result.binary_results_for_llm is not None
        assert len(rpc_result.binary_results_for_llm) == 1
        assert rpc_result.binary_results_for_llm[0].data == "base64data"
        assert rpc_result.binary_results_for_llm[0].mime_type == "image/png"
        assert rpc_result.binary_results_for_llm[0].type.value == "image"
        assert rpc_result.binary_results_for_llm[0].description == "screenshot.png"

    def test_omits_binary_results_when_none(self):
        tool_result = ToolResult(text_result_for_llm="done")
        rpc_result = tool_result_to_external_tool_text_result_for_llm(tool_result)
        assert rpc_result.binary_results_for_llm is None
        assert rpc_result.session_log is None

    def test_forwards_tool_references(self):
        tool_result = ToolResult(
            text_result_for_llm="found tools",
            result_type="success",
            tool_references=["get_weather", "check_status"],
        )
        rpc_result = tool_result_to_external_tool_text_result_for_llm(tool_result)
        assert rpc_result.tool_references == ["get_weather", "check_status"]