wrolpi/modules/ai/test/test_agent.py

"""Tests for tool-definition generation and tool execution."""
import json
from unittest import mock
from urllib.parse import urlencode

import pytest

from modules.ai import agent
from modules.ai.modes import MODES


@pytest.fixture
async def tool_definitions(async_client):
    # Any request finalizes the router, which the spec builder requires.
    await async_client.get('/api/ai/zims')
    return agent.get_ai_tool_definitions(refresh=True)


@pytest.mark.asyncio
async def test_tool_definitions_from_spec(tool_definitions):
    """Tool definitions are generated from the blueprint's own OpenAPI operations."""
    tools = tool_definitions
    # Every tool any mode allowlists must exist — this guards allowlists against drift.
    for mode_name, mode in MODES.items():
        for tool_name in mode['tools']:
            assert tool_name in tools, f'Mode {mode_name} allowlists missing tool {tool_name}'

    # Every tool has a model-facing description.
    for name, tool in tools.items():
        assert tool.description, f'Tool {name} has no description'

    # Manage endpoints are never tools.
    assert not any(t.path.startswith('/api/ai/manage') for t in tools.values())
    assert 'manage_catalog' not in tools and 'manage_settings' not in tools
    # The chat endpoint itself is not a tool.
    assert not any(t.path.startswith('/api/chat') for t in tools.values())

    # Body, query, and path parameters all appear in the flat schema.
    schema = tools['search_videos'].parameters_schema()
    assert schema['properties']['search_str'] == {'type': 'string'}
    assert schema['properties']['tag_names'] == {'type': 'array', 'items': {'type': 'string'}}
    schema = tools['get_video_captions'].parameters_schema()
    assert schema['required'] == ['file_group_id']
    assert schema['properties']['offset'] == {'type': 'integer'}

    # OpenAI tool format.
    openai_tool = tools['search_all'].as_openai_tool()
    assert openai_tool['type'] == 'function'
    assert openai_tool['function']['name'] == 'search_all'


def _endpoint_shim(async_client):
    async def shim(method, path, query, body):
        url = f'{path}?{urlencode(query)}' if query else path
        if method == 'post':
            request, response = await async_client.post(url, content=json.dumps(body or {}))
        else:
            request, response = await async_client.get(url)
        return response.status_code, response.json

    return mock.patch('modules.ai.agent.call_ai_endpoint', side_effect=shim)


@pytest.mark.asyncio
async def test_execute_tool(async_client, tool_definitions, video_factory, test_session):
    """Tools execute against the blueprint; every failure is an instructive error string."""
    video_factory(title='pressure canning')
    test_session.commit()
    research_tools = MODES['research']['tools']

    with _endpoint_shim(async_client):
        # A real search through the blueprint.
        success, result = await agent.execute_tool(research_tools, 'search_videos',
                                                   dict(search_str='canning'))
        assert success is True
        result = json.loads(result)
        assert result['total'] == 1

        # Unknown tool.
        success, error = await agent.execute_tool(research_tools, 'delete_everything', {})
        assert success is False and 'Unknown tool' in error

        # Mode allowlist: research mode cannot read service logs even though the tool exists.
        success, error = await agent.execute_tool(research_tools, 'get_service_logs',
                                                  dict(name='wrolpi-api'))
        assert success is False and 'Unknown tool' in error

        # Unknown argument names the accepted ones.
        success, error = await agent.execute_tool(research_tools, 'search_videos', dict(query='x'))
        assert success is False and 'accepts' in error and 'search_str' in error

        # Missing required path parameter.
        success, error = await agent.execute_tool(research_tools, 'get_video_captions', {})
        assert success is False and 'Missing required' in error

        # Wrong type.
        success, error = await agent.execute_tool(research_tools, 'get_video_captions',
                                                  dict(file_group_id='banana'))
        assert success is False and 'must be a integer' in error

        # An endpoint error comes back as a message, not an exception.
        success, error = await agent.execute_tool(research_tools, 'get_video_captions',
                                                  dict(file_group_id=123456))
        assert success is False and 'error' in error.lower()


@pytest.mark.asyncio
async def test_execute_tool_path_traversal(async_client, tool_definitions):
    """A crafted path parameter cannot escape the tool's endpoint (allowlist bypass)."""
    help_tools = MODES['help']['tools']
    with _endpoint_shim(async_client):
        for slug in ('../../manage/catalog', '../manage/settings', '/etc/passwd',
                     'a/../../files/refresh', 'a//b', 'a/./b'):
            success, error = await agent.execute_tool(help_tools, 'get_help_doc', dict(slug=slug))
            assert success is False, f'traversal not refused: {slug}'
            assert 'traversal' in error or 'invalid path' in error

        # Legitimate nested slugs keep working (404 from the endpoint is fine, no traversal error).
        success, error = await agent.execute_tool(help_tools, 'get_help_doc', dict(slug='system/logs'))
        assert 'traversal' not in error and 'invalid path' not in (error or '')


@pytest.mark.asyncio
async def test_execute_tool_result_truncated(async_client, tool_definitions, archive_factory, test_session):
    """A huge tool result is capped with a paging hint."""
    archive_factory(domain='example.com', title='big page', contents='word ' * 5_000)
    test_session.commit()
    from modules.archive.models import Archive
    archive = test_session.query(Archive).one()

    with _endpoint_shim(async_client):
        success, result = await agent.execute_tool(MODES['research']['tools'], 'get_archive_text',
                                                   dict(file_group_id=archive.file_group_id))
        assert success is True
        assert len(result) <= agent.TOOL_RESULT_MAX_CHARS + 100