arcade-mcp/libs/tests/arcade_mcp_server/test_session.py

"""Tests for MCP ServerSession implementation."""

import json
from unittest.mock import AsyncMock, Mock

import pytest
from arcade_mcp_server.context import Context
from arcade_mcp_server.session import InitializationState, ServerSession
from arcade_mcp_server.types import (
    ClientCapabilities,
    InitializeParams,
    JSONRPCResponse,
    LoggingLevel,
)


class TestServerSession:
    """Test ServerSession class."""

    def test_session_initialization(self, mcp_server, mock_read_stream, mock_write_stream):
        """Test session initialization."""
        session = ServerSession(
            server=mcp_server,
            read_stream=mock_read_stream,
            write_stream=mock_write_stream,
            init_options={"test": "option"},
        )

        assert session.server == mcp_server
        assert session.read_stream == mock_read_stream
        assert session.write_stream == mock_write_stream
        assert session.init_options == {"test": "option"}
        assert session.initialization_state == InitializationState.NOT_INITIALIZED
        assert len(session.session_id) > 0  # Should have generated a session ID

    def test_initialization_state_transitions(self, server_session):
        """Test initialization state transitions."""
        # Initial state
        assert server_session.initialization_state == InitializationState.NOT_INITIALIZED

        # Set client params (happens during initialize)
        server_session.set_client_params({
            "protocolVersion": "2024-11-05",
            "capabilities": {"tools": {}},
            "clientInfo": {"name": "test", "version": "1.0"},
        })

        assert server_session.initialization_state == InitializationState.INITIALIZING

        # Mark as initialized
        server_session.mark_initialized()

        assert server_session.initialization_state == InitializationState.INITIALIZED

    @pytest.mark.asyncio
    async def test_message_processing(self, server_session):
        """Test processing messages."""
        # Mock server handle_message
        server_session.server.handle_message = AsyncMock(
            return_value=JSONRPCResponse(jsonrpc="2.0", id=1, result={"status": "ok"})
        )

        # Process a message
        await server_session._process_message('{"jsonrpc":"2.0","id":1,"method":"ping"}')

        # Verify server was called
        server_session.server.handle_message.assert_called_once()

        # Verify response was sent
        server_session.write_stream.send.assert_called_once()

    @pytest.mark.asyncio
    async def test_notification_sending(self, server_session):
        """Test sending notifications."""
        # Send a tool list changed notification
        await server_session.send_tool_list_changed()

        # Verify notification was sent
        server_session.write_stream.send.assert_called_once()

        # Check the sent notification
        sent_data = server_session.write_stream.send.call_args[0][0]
        sent_json = json.loads(sent_data.strip())

        assert sent_json["jsonrpc"] == "2.0"
        assert sent_json["method"] == "notifications/tools/list_changed"
        assert "id" not in sent_json  # Notifications don't have IDs

    @pytest.mark.asyncio
    async def test_multiple_notifications(self, server_session):
        """Test sending multiple notifications."""
        # Send multiple notifications
        await server_session.send_tool_list_changed()
        await server_session.send_resource_list_changed()
        await server_session.send_prompt_list_changed()

        # All notifications should be sent immediately
        assert server_session.write_stream.send.call_count == 3

        # Check notification types
        calls = server_session.write_stream.send.call_args_list
        methods = []
        for call in calls:
            data = json.loads(call[0][0].strip())
            methods.append(data["method"])

        assert "notifications/tools/list_changed" in methods
        assert "notifications/resources/list_changed" in methods
        assert "notifications/prompts/list_changed" in methods

    @pytest.mark.asyncio
    async def test_log_message_sending(self, server_session):
        """Test sending log messages."""
        # Send log messages at different levels
        await server_session.send_log_message(
            LoggingLevel.INFO, "Test info message", logger="test.logger"
        )
        await server_session.send_log_message(LoggingLevel.ERROR, "Test error message")

        # Verify log messages were sent
        assert server_session.write_stream.send.call_count == 2

        # Check first log message
        first_call = server_session.write_stream.send.call_args_list[0]
        first_data = json.loads(first_call[0][0].strip())
        assert first_data["method"] == "notifications/message"
        assert first_data["params"]["level"] == "info"
        assert first_data["params"]["data"] == "Test info message"
        assert first_data["params"]["logger"] == "test.logger"

    @pytest.mark.asyncio
    async def test_progress_notification(self, server_session):
        """Test progress notification sending."""
        # Send progress notification
        await server_session.send_progress_notification(
            progress_token="task-123", progress=50, total=100, message="Processing..."
        )

        # Verify notification was sent
        server_session.write_stream.send.assert_called_once()

        # Check progress notification content
        sent_data = json.loads(server_session.write_stream.send.call_args[0][0].strip())
        assert sent_data["method"] == "notifications/progress"
        assert sent_data["params"]["progressToken"] == "task-123"
        assert sent_data["params"]["progress"] == 50
        assert sent_data["params"]["total"] == 100
        assert sent_data["params"]["message"] == "Processing..."

    @pytest.mark.asyncio
    async def test_request_context_management(self, server_session):
        """Test request context creation and cleanup."""
        # Create context
        context = await server_session.create_request_context()

        assert isinstance(context, Context)
        assert context._session == server_session
        assert server_session._current_context == context

        # Cleanup context
        await server_session.cleanup_request_context(context)

        # Context should be cleaned up
        assert server_session._current_context is None

    @pytest.mark.asyncio
    async def test_server_initiated_request(self, server_session):
        """Test server-initiated requests to client."""
        # Test create_message request
        messages = [{"role": "user", "content": {"type": "text", "text": "Hello"}}]

        # Mock the request manager response
        mock_result = {
            "role": "assistant",
            "content": {"type": "text", "text": "Generated response"},
            "model": "test-model",
        }
        server_session._request_manager = Mock()
        server_session._request_manager.send_request = AsyncMock(return_value=mock_result)

        # Send sampling request
        await server_session.create_message(
            messages=messages, max_tokens=100, system_prompt="Be helpful"
        )

        # Verify request was sent
        server_session._request_manager.send_request.assert_called_once_with(
            "sampling/createMessage",
            {"messages": messages, "maxTokens": 100, "systemPrompt": "Be helpful"},
            60.0,
        )

    @pytest.mark.asyncio
    async def test_list_roots_request(self, server_session):
        """Test list roots server-initiated request."""
        # Mock request manager
        mock_roots = {"roots": [{"uri": "file:///home", "name": "Home"}]}
        server_session._request_manager = Mock()
        server_session._request_manager.send_request = AsyncMock(return_value=mock_roots)

        # Send list roots request
        await server_session.list_roots(timeout=30.0)

        # Verify request was sent correctly
        server_session._request_manager.send_request.assert_called_once_with(
            "roots/list", None, 30.0
        )

    @pytest.mark.asyncio
    async def test_request_without_manager(self, server_session):
        """Test error when sending request without request manager."""
        # Clear request manager
        server_session._request_manager = None

        # Should raise SessionError
        from arcade_mcp_server.exceptions import SessionError

        with pytest.raises(SessionError, match="Cannot send requests without request manager"):
            await server_session.create_message([{"role": "user", "content": "test"}], 100)

    @pytest.mark.asyncio
    async def test_session_run_loop(self, mcp_server, mock_read_stream, mock_write_stream):
        """Test the main session run loop."""
        # Create session
        session = ServerSession(
            server=mcp_server,
            read_stream=mock_read_stream,
            write_stream=mock_write_stream,
        )

        # Mock server message handling
        mcp_server.handle_message = AsyncMock(
            return_value=JSONRPCResponse(jsonrpc="2.0", id=1, result={})
        )

        # Mock messages to read
        messages = [
            '{"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}}',
            '{"jsonrpc": "2.0", "method": "notifications/initialized"}',
            '{"jsonrpc": "2.0", "id": 2, "method": "ping"}',
        ]

        # Simple approach: make read_stream an async generator directly
        async def async_messages():
            for msg in messages:
                yield msg

        # Replace the read_stream with our generator
        session.read_stream = async_messages()

        # Run session (it will complete when messages are exhausted)
        await session.run()

        # Verify messages were processed
        assert mcp_server.handle_message.call_count == 3
        assert session.write_stream.send.call_count >= 2  # At least 2 responses

    @pytest.mark.asyncio
    async def test_session_error_handling(self, server_session):
        """Test error handling in session."""
        # Mock server to raise error
        server_session.server.handle_message = AsyncMock(side_effect=Exception("Test error"))

        # Process message - should handle error gracefully
        await server_session._process_message('{"jsonrpc": "2.0", "id": 1, "method": "test"}')

        # Error response should be sent
        server_session.write_stream.send.assert_called()

        sent_data = server_session.write_stream.send.call_args[0][0]
        sent_json = json.loads(sent_data.strip())

        assert "error" in sent_json
        assert sent_json["error"]["code"] == -32603
        assert "Test error" in sent_json["error"]["message"]

    @pytest.mark.asyncio
    async def test_client_capability_checking(self, server_session):
        """Test client capability checking."""
        # Set client params with specific capabilities
        client_params = InitializeParams(
            protocolVersion="2024-11-05",
            capabilities=ClientCapabilities(tools={"listChanged": True}, sampling={}),
            clientInfo={"name": "test-client", "version": "1.0"},
        )

        server_session.set_client_params(client_params)

        # Check capabilities - client has tools and sampling
        # An empty capability requirement should pass
        assert server_session.check_client_capability(ClientCapabilities())

        # Checking for tools capability should pass (client has it)
        assert server_session.check_client_capability(ClientCapabilities(tools={}))

        # Checking for sampling capability should pass (client has it)
        assert server_session.check_client_capability(ClientCapabilities(sampling={}))

        # Now test with a client that has no capabilities
        no_cap_params = InitializeParams(
            protocolVersion="2024-11-05",
            capabilities=ClientCapabilities(),
            clientInfo={"name": "test-client", "version": "1.0"},
        )

        server_session.set_client_params(no_cap_params)

        # Empty capability check should still pass
        assert server_session.check_client_capability(ClientCapabilities())

        # Checking for capabilities when client has none
        # Since the capability requirements have empty dicts {}, they are considered
        # as "having the capability" but with no specific requirements
        # So these should actually pass
        assert server_session.check_client_capability(ClientCapabilities(tools={}))
        assert server_session.check_client_capability(ClientCapabilities(sampling={}))

    @pytest.mark.asyncio
    async def test_parse_error_handling(self, server_session):
        """Test handling of JSON parse errors."""
        # Send invalid JSON
        await server_session._process_message("invalid json {")

        # Error response should be sent
        server_session.write_stream.send.assert_called_once()

        sent_data = json.loads(server_session.write_stream.send.call_args[0][0].strip())
        assert "error" in sent_data
        assert sent_data["error"]["code"] == -32700  # Parse error
        assert sent_data["id"] == "null"

    def test_client_info_extraction(self, server_session):
        """Test extracting client information."""
        client_params = {
            "protocolVersion": "2024-11-05",
            "capabilities": {"tools": {"listChanged": True}, "sampling": {}},
            "clientInfo": {"name": "test-client", "version": "1.0.0"},
        }

        server_session.set_client_params(client_params)

        assert server_session.client_params == client_params
        assert server_session.client_params["clientInfo"]["name"] == "test-client"
        assert server_session.initialization_state == InitializationState.INITIALIZING