ServerKit/backend/tests/test_ai_assistant.py

"""Tests for the core AI assistant primitive (powered by Prompture)."""
import pytest

from prompture.agents.live_events import (
    TextDelta, ToolResult, ToolUseStart, TurnComplete,
)

from app.services import ai_service
from app.services.ai_tool_registry import ai_tool_registry


class FakeUser:
    def __init__(self, role='admin', perms=None):
        self.role = role
        self.id = 1
        self._perms = perms or {}

    def has_permission(self, feature, level='read'):
        if self.role == 'admin':
            return True
        return self._perms.get((feature, level), False)


# --------------------------------------------------------------------------
# LiveEvent -> SSE frame mapping
# --------------------------------------------------------------------------
def test_live_event_to_frame_text_delta():
    name, data = ai_service.live_event_to_frame(TextDelta(text='hello'))
    assert name == 'text_delta'
    assert data == {'text': 'hello'}


def test_live_event_to_frame_tool_use_start():
    name, data = ai_service.live_event_to_frame(ToolUseStart(id='t1', name='core__list_apps'))
    assert name == 'tool_use_start'
    assert data == {'id': 't1', 'name': 'core__list_apps'}


def test_live_event_to_frame_tool_result():
    name, data = ai_service.live_event_to_frame(
        ToolResult(id='t1', name='core__x', output='ok', is_error=False))
    assert name == 'tool_result'
    assert data['id'] == 't1' and data['output'] == 'ok' and data['is_error'] is False


def test_live_event_to_frame_turn_complete():
    name, data = ai_service.live_event_to_frame(TurnComplete(usage={'cost': 0.12}))
    assert name == 'turn_complete'
    assert data['usage'] == {'cost': 0.12}


# --------------------------------------------------------------------------
# Tool registry: register / unregister / reload + RBAC + namespacing
# --------------------------------------------------------------------------
def test_registry_register_and_unregister():
    def dbl(x: int) -> int:
        """Double a number.

        Args:
            x: the number to double.
        """
        return x * 2

    ai_tool_registry.register(name='dbl', func=dbl, plugin_slug='tpA',
                              rbac_feature='docker', rbac_level='read')
    d = ai_tool_registry.get('tpA__dbl')
    assert d is not None
    assert d.qualified_name == 'tpA__dbl'
    assert d.is_write is False
    assert d.parameters['properties']['x']['type'] == 'integer'

    ai_tool_registry.unregister_plugin('tpA')
    assert ai_tool_registry.get('tpA__dbl') is None


def test_registry_rbac_and_mode_filter():
    def peek() -> str:
        """Peek at something secret."""
        return 'ok'

    ai_tool_registry.register(name='peek', func=peek, plugin_slug='tpB',
                              rbac_feature='security', rbac_level='read')
    try:
        admin = FakeUser(role='admin')
        viewer = FakeUser(role='viewer', perms={('security', 'read'): False})

        admin_names = [d.qualified_name for d in ai_tool_registry.list_for(admin, 'assistant')]
        viewer_names = [d.qualified_name for d in ai_tool_registry.list_for(viewer, 'assistant')]
        assert 'tpB__peek' in admin_names
        assert 'tpB__peek' not in viewer_names

        # Simple mode never offers tools.
        assert ai_tool_registry.list_for(admin, 'simple') == []
    finally:
        ai_tool_registry.unregister_plugin('tpB')


def test_plugin_binder_namespaces_and_flags():
    from app.plugins_sdk.ai import PluginToolBinder
    binder = PluginToolBinder('demo')

    @binder.tool(rbac_feature='git', is_write=True)
    def delete_branch(repo: str, branch: str) -> dict:
        """Delete a branch.

        Args:
            repo: repository slug.
            branch: branch name.
        """
        return {}

    try:
        d = ai_tool_registry.get('demo__delete_branch')
        assert d is not None
        assert d.plugin_slug == 'demo'
        assert d.is_write is True
        assert d.rbac_level == 'write'
    finally:
        ai_tool_registry.unregister_plugin('demo')


# --------------------------------------------------------------------------
# Provider environment construction (Prompture ProviderEnvironment)
# --------------------------------------------------------------------------
def test_build_provider_env_openai_key(monkeypatch):
    monkeypatch.setattr(ai_service, '_setting',
                        lambda k, d=None: {'ai_provider': 'openai', 'ai_endpoint': ''}.get(k, d))
    monkeypatch.setattr(ai_service, '_decrypted_key', lambda: 'sk-test-123')
    env = ai_service.build_provider_env()
    assert env.openai_api_key == 'sk-test-123'
    assert env.ollama_endpoint is None


def test_build_provider_env_ollama_endpoint(monkeypatch):
    monkeypatch.setattr(ai_service, '_setting',
                        lambda k, d=None: {'ai_provider': 'ollama', 'ai_endpoint': 'http://h:11434'}.get(k, d))
    monkeypatch.setattr(ai_service, '_decrypted_key', lambda: None)
    env = ai_service.build_provider_env()
    assert env.ollama_endpoint == 'http://h:11434'
    assert env.openai_api_key is None


def test_summarize_action_includes_params():
    class D:
        name = 'restart_docker_container'
        description = 'Restart a Docker container. STATE-CHANGING'

    summary = ai_service.summarize_action(D(), {'container_id': 'nginx'})
    assert 'container_id=nginx' in summary


# --------------------------------------------------------------------------
# API: RBAC gating + secret never leaked
# --------------------------------------------------------------------------
def test_status_requires_auth_and_reports_configured(client, auth_headers):
    r = client.get('/api/v1/ai/status', headers=auth_headers)
    assert r.status_code == 200
    body = r.get_json()
    assert 'configured' in body and 'enabled' in body


def test_settings_get_never_returns_api_key(client, auth_headers, app):
    from app.services.settings_service import SettingsService
    with app.app_context():
        SettingsService.set('ai_api_key_encrypted', 'ENCRYPTED_SECRET_XYZ')

    r = client.get('/api/v1/ai/settings', headers=auth_headers)
    assert r.status_code == 200
    body = r.get_json()
    assert body['api_key_set'] is True
    assert 'api_key' not in body
    assert 'ENCRYPTED_SECRET_XYZ' not in r.get_data(as_text=True)


def test_settings_put_requires_admin(client, app):
    from app import db
    from app.models import User
    from flask_jwt_extended import create_access_token
    from werkzeug.security import generate_password_hash

    with app.app_context():
        viewer = User(
            email='viewer@test.local', username='viewer_ai',
            password_hash=generate_password_hash('x'),
            role=User.ROLE_VIEWER, is_active=True,
        )
        db.session.add(viewer)
        db.session.commit()
        token = create_access_token(identity=viewer.id)

    r = client.put('/api/v1/ai/settings',
                   headers={'Authorization': f'Bearer {token}'},
                   json={'provider': 'openai'})
    assert r.status_code == 403