test_tools_e2e.py

"""E2E Tools Tests"""

import os

import pytest
from pydantic import BaseModel, Field

from copilot import define_tool
from copilot.generated.rpc import PermissionDecisionApproveOnce, PermissionDecisionReject
from copilot.session import PermissionHandler, PermissionNoResult
from copilot.tools import Tool, ToolInvocation, ToolResult

from .testharness import E2ETestContext, get_final_assistant_message

pytestmark = pytest.mark.asyncio(loop_scope="module")


class TestTools:
    async def test_invokes_built_in_tools(self, ctx: E2ETestContext):
        readme_path = os.path.join(ctx.work_dir, "README.md")
        with open(readme_path, "w") as f:
            f.write("# ELIZA, the only chatbot you'll ever need")

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all
        )

        await session.send("What's the first line of README.md in this directory?")
        assistant_message = await get_final_assistant_message(session)
        assert "ELIZA" in assistant_message.data.content

    async def test_invokes_custom_tool(self, ctx: E2ETestContext):
        class EncryptParams(BaseModel):
            input: str = Field(description="String to encrypt")

        @define_tool("encrypt_string", description="Encrypts a string")
        def encrypt_string(params: EncryptParams, invocation: ToolInvocation) -> str:
            return params.input.upper()

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all, tools=[encrypt_string]
        )

        await session.send("Use encrypt_string to encrypt this string: Hello")
        assistant_message = await get_final_assistant_message(session)
        assert "HELLO" in assistant_message.data.content

    async def test_handles_tool_calling_errors(self, ctx: E2ETestContext):
        @define_tool("get_user_location", description="Gets the user's location")
        def get_user_location() -> str:
            raise Exception("Melbourne")

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all, tools=[get_user_location]
        )

        await session.send("What is my location? If you can't find out, just say 'unknown'.")
        answer = await get_final_assistant_message(session)

        # Check the underlying traffic
        traffic = await ctx.get_exchanges()
        last_conversation = traffic[-1]

        tool_calls = []
        for msg in last_conversation["request"]["messages"]:
            if msg.get("role") == "assistant" and "tool_calls" in msg:
                tool_calls.extend(msg["tool_calls"])

        assert len(tool_calls) == 1
        tool_call = tool_calls[0]
        assert tool_call["type"] == "function"
        assert tool_call["function"]["name"] == "get_user_location"

        tool_results = [
            msg for msg in last_conversation["request"]["messages"] if msg.get("role") == "tool"
        ]
        assert len(tool_results) == 1
        tool_result = tool_results[0]
        assert tool_result["tool_call_id"] == tool_call["id"]

        # The error message "Melbourne" should NOT be exposed to the LLM
        assert "Melbourne" not in tool_result["content"]

        # The assistant should not see the exception information
        assert "Melbourne" not in (answer.data.content or "")
        assert "unknown" in (answer.data.content or "").lower()

    async def test_can_receive_and_return_complex_types(self, ctx: E2ETestContext):
        class DbQuery(BaseModel):
            table: str
            ids: list[int]
            sortAscending: bool

        class DbQueryParams(BaseModel):
            query: DbQuery

        class City(BaseModel):
            countryId: int
            cityName: str
            population: int

        expected_session_id = None

        @define_tool("db_query", description="Performs a database query")
        def db_query(params: DbQueryParams, invocation: ToolInvocation) -> list[City]:
            assert params.query.table == "cities"
            assert params.query.ids == [12, 19]
            assert params.query.sortAscending is True
            assert invocation.session_id == expected_session_id

            return [
                City(countryId=19, cityName="Passos", population=135460),
                City(countryId=12, cityName="San Lorenzo", population=204356),
            ]

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all, tools=[db_query]
        )
        expected_session_id = session.session_id

        await session.send(
            "Perform a DB query for the 'cities' table using IDs 12 and 19, "
            "sorting ascending. Reply only with lines of the form: [cityname] [population]"
        )

        assistant_message = await get_final_assistant_message(session)
        response_content = assistant_message.data.content or ""

        assert response_content != ""
        assert "Passos" in response_content
        assert "San Lorenzo" in response_content
        assert "135460" in response_content.replace(",", "")
        assert "204356" in response_content.replace(",", "")

    async def test_skippermission_sent_in_tool_definition(self, ctx: E2ETestContext):
        class LookupParams(BaseModel):
            id: str = Field(description="ID to look up")

        @define_tool(
            "safe_lookup",
            description="A safe lookup that skips permission",
            skip_permission=True,
        )
        def safe_lookup(params: LookupParams, invocation: ToolInvocation) -> str:
            return f"RESULT: {params.id}"

        did_run_permission_request = False

        def tracking_handler(request, invocation):
            nonlocal did_run_permission_request
            did_run_permission_request = True
            return PermissionNoResult()

        session = await ctx.client.create_session(
            on_permission_request=tracking_handler, tools=[safe_lookup]
        )

        await session.send("Use safe_lookup to look up 'test123'")
        assistant_message = await get_final_assistant_message(session)
        assert "RESULT: test123" in assistant_message.data.content
        assert not did_run_permission_request

    async def test_overrides_built_in_tool_with_custom_tool(self, ctx: E2ETestContext):
        class GrepParams(BaseModel):
            query: str = Field(description="Search query")

        @define_tool(
            "grep",
            description="A custom grep implementation that overrides the built-in",
            overrides_built_in_tool=True,
        )
        def custom_grep(params: GrepParams, invocation: ToolInvocation) -> str:
            return f"CUSTOM_GREP_RESULT: {params.query}"

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all, tools=[custom_grep]
        )

        await session.send("Use grep to search for the word 'hello'")
        assistant_message = await get_final_assistant_message(session)
        assert "CUSTOM_GREP_RESULT" in assistant_message.data.content

    async def test_invokes_custom_tool_with_permission_handler(self, ctx: E2ETestContext):
        class EncryptParams(BaseModel):
            input: str = Field(description="String to encrypt")

        @define_tool("encrypt_string", description="Encrypts a string")
        def encrypt_string(params: EncryptParams, invocation: ToolInvocation) -> str:
            return params.input.upper()

        permission_requests = []

        def on_permission_request(request, invocation):
            permission_requests.append(request)
            return PermissionDecisionApproveOnce()

        session = await ctx.client.create_session(
            on_permission_request=on_permission_request, tools=[encrypt_string]
        )

        await session.send("Use encrypt_string to encrypt this string: Hello")
        assistant_message = await get_final_assistant_message(session)
        assert "HELLO" in assistant_message.data.content

        # Should have received a custom-tool permission request
        custom_tool_requests = [r for r in permission_requests if r.kind == "custom-tool"]
        assert len(custom_tool_requests) > 0
        assert custom_tool_requests[0].tool_name == "encrypt_string"

    async def test_denies_custom_tool_when_permission_denied(self, ctx: E2ETestContext):
        tool_handler_called = False

        class EncryptParams(BaseModel):
            input: str = Field(description="String to encrypt")

        @define_tool("encrypt_string", description="Encrypts a string")
        def encrypt_string(params: EncryptParams, invocation: ToolInvocation) -> str:
            nonlocal tool_handler_called
            tool_handler_called = True
            return params.input.upper()

        def on_permission_request(request, invocation):
            return PermissionDecisionReject()

        session = await ctx.client.create_session(
            on_permission_request=on_permission_request, tools=[encrypt_string]
        )

        await session.send("Use encrypt_string to encrypt this string: Hello")
        await get_final_assistant_message(session)

        # The tool handler should NOT have been called since permission was denied
        assert not tool_handler_called

    async def test_should_execute_multiple_custom_tools_in_parallel_single_turn(
        self, ctx: E2ETestContext
    ):
        """Multiple custom tools invoked in parallel in the same turn."""
        import asyncio

        city_called: asyncio.Future = asyncio.get_event_loop().create_future()
        country_called: asyncio.Future = asyncio.get_event_loop().create_future()

        def lookup_city(invocation: ToolInvocation) -> ToolResult:
            city = (invocation.arguments or {}).get("city", "")
            if not city_called.done():
                city_called.set_result(city)
            return ToolResult(text_result_for_llm=f"CITY_{city.upper()}")

        def lookup_country(invocation: ToolInvocation) -> ToolResult:
            country = (invocation.arguments or {}).get("country", "")
            if not country_called.done():
                country_called.set_result(country)
            return ToolResult(text_result_for_llm=f"COUNTRY_{country.upper()}")

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
            tools=[
                Tool(
                    name="lookup_city",
                    description="Looks up city information",
                    parameters={
                        "type": "object",
                        "properties": {"city": {"type": "string", "description": "City name"}},
                        "required": ["city"],
                    },
                    handler=lookup_city,
                ),
                Tool(
                    name="lookup_country",
                    description="Looks up country information",
                    parameters={
                        "type": "object",
                        "properties": {
                            "country": {"type": "string", "description": "Country name"}
                        },
                        "required": ["country"],
                    },
                    handler=lookup_country,
                ),
            ],
        )

        try:
            await session.send(
                "Use lookup_city with 'Paris' and lookup_country with 'France' at the same time,"
                " then combine both results in your reply."
            )

            city_result = await asyncio.wait_for(city_called, timeout=60.0)
            country_result = await asyncio.wait_for(country_called, timeout=60.0)
            assert city_result == "Paris"
            assert country_result == "France"

            assistant_message = await get_final_assistant_message(session, timeout=60.0)
            assert assistant_message is not None
            content = assistant_message.data.content or ""
            assert "CITY_PARIS" in content
            assert "COUNTRY_FRANCE" in content
        finally:
            await session.disconnect()

    async def test_should_respect_availabletools_and_excludedtools_combined(
        self, ctx: E2ETestContext
    ):
        """excluded_tools takes precedence over available_tools."""
        excluded_tool_called = False

        def allowed_handler(invocation: ToolInvocation) -> ToolResult:
            input_val = (invocation.arguments or {}).get("input", "")
            return ToolResult(text_result_for_llm=f"ALLOWED_{input_val.upper()}")

        def excluded_handler(invocation: ToolInvocation) -> ToolResult:
            nonlocal excluded_tool_called
            excluded_tool_called = True
            input_val = (invocation.arguments or {}).get("input", "")
            return ToolResult(text_result_for_llm=f"EXCLUDED_{input_val.upper()}")

        session = await ctx.client.create_session(
            on_permission_request=PermissionHandler.approve_all,
            tools=[
                Tool(
                    name="allowed_tool",
                    description="An allowed tool",
                    parameters={
                        "type": "object",
                        "properties": {"input": {"type": "string", "description": "Input value"}},
                        "required": ["input"],
                    },
                    handler=allowed_handler,
                ),
                Tool(
                    name="excluded_tool",
                    description="A tool that should be excluded",
                    parameters={
                        "type": "object",
                        "properties": {"input": {"type": "string", "description": "Input value"}},
                        "required": ["input"],
                    },
                    handler=excluded_handler,
                ),
            ],
            available_tools=["allowed_tool", "excluded_tool"],
            excluded_tools=["excluded_tool"],
        )

        try:
            result = await session.send_and_wait(
                "Use the allowed_tool with input 'test'. Do NOT use excluded_tool.",
                timeout=60.0,
            )
            assert result is not None
            assert "ALLOWED_TEST" in (result.data.content or "")
            assert not excluded_tool_called, "Excluded tool should not have been called"
        finally:
            await session.disconnect()