atlas-chat 0.1.0__py3-none-any.whl

atlas/tests/test_elicitation_manager.py

@@ -0,0 +1,408 @@
+"""Tests for the elicitation manager."""
+
+import asyncio
+
+import pytest
+
+from atlas.application.chat.elicitation_manager import ElicitationManager, ElicitationRequest, get_elicitation_manager
+
+
+class TestElicitationRequest:
+    """Test ElicitationRequest class."""
+
+    @pytest.mark.asyncio
+    async def test_create_elicitation_request(self):
+        """Test creating an elicitation request."""
+        request = ElicitationRequest(
+            elicitation_id="elicit_123",
+            tool_call_id="tool_456",
+            tool_name="test_tool",
+            message="Please provide your name",
+            response_schema={"type": "object", "properties": {"value": {"type": "string"}}}
+        )
+        assert request.elicitation_id == "elicit_123"
+        assert request.tool_call_id == "tool_456"
+        assert request.tool_name == "test_tool"
+        assert request.message == "Please provide your name"
+        assert "properties" in request.response_schema
+
+    @pytest.mark.asyncio
+    async def test_wait_for_response_accept(self):
+        """Test waiting for an accept response."""
+        request = ElicitationRequest(
+            elicitation_id="elicit_123",
+            tool_call_id="tool_456",
+            tool_name="test_tool",
+            message="Enter your name",
+            response_schema={"type": "object"}
+        )
+
+        # Simulate setting a response
+        response_data = {"action": "accept", "data": {"name": "John"}}
+        request.future.set_result(response_data)
+
+        # Wait for the response (should be immediate since we already set it)
+        response = await request.wait_for_response(timeout=1.0)
+
+        assert response["action"] == "accept"
+        assert response["data"] == {"name": "John"}
+
+    @pytest.mark.asyncio
+    async def test_wait_for_response_decline(self):
+        """Test waiting for a decline response."""
+        request = ElicitationRequest(
+            elicitation_id="elicit_123",
+            tool_call_id="tool_456",
+            tool_name="test_tool",
+            message="Enter your name",
+            response_schema={"type": "object"}
+        )
+
+        # Simulate decline response
+        response_data = {"action": "decline", "data": None}
+        request.future.set_result(response_data)
+
+        response = await request.wait_for_response(timeout=1.0)
+
+        assert response["action"] == "decline"
+        assert response["data"] is None
+
+    @pytest.mark.asyncio
+    async def test_wait_for_response_cancel(self):
+        """Test waiting for a cancel response."""
+        request = ElicitationRequest(
+            elicitation_id="elicit_123",
+            tool_call_id="tool_456",
+            tool_name="test_tool",
+            message="Enter your name",
+            response_schema={"type": "object"}
+        )
+
+        # Simulate cancel response
+        response_data = {"action": "cancel", "data": None}
+        request.future.set_result(response_data)
+
+        response = await request.wait_for_response(timeout=1.0)
+
+        assert response["action"] == "cancel"
+        assert response["data"] is None
+
+    @pytest.mark.asyncio
+    async def test_timeout(self):
+        """Test that timeout works correctly."""
+        request = ElicitationRequest(
+            elicitation_id="elicit_123",
+            tool_call_id="tool_456",
+            tool_name="test_tool",
+            message="Enter your name",
+            response_schema={"type": "object"}
+        )
+
+        # Should timeout since we don't set a response
+        with pytest.raises(asyncio.TimeoutError):
+            await request.wait_for_response(timeout=0.1)
+
+
+class TestElicitationManager:
+    """Test ElicitationManager class."""
+
+    @pytest.mark.asyncio
+    async def test_create_elicitation_request(self):
+        """Test creating an elicitation request via manager."""
+        manager = ElicitationManager()
+        manager.create_elicitation_request(
+            elicitation_id="elicit_123",
+            tool_call_id="tool_456",
+            tool_name="test_tool",
+            message="Enter your name",
+            response_schema={"type": "object"}
+        )
+
+        assert "elicit_123" in manager.get_all_pending_requests()
+
+    @pytest.mark.asyncio
+    async def test_handle_elicitation_response_accept(self):
+        """Test handling an accept response."""
+        manager = ElicitationManager()
+        manager.create_elicitation_request(
+            elicitation_id="elicit_123",
+            tool_call_id="tool_456",
+            tool_name="test_tool",
+            message="Enter your name",
+            response_schema={"type": "object"}
+        )
+
+        # Handle the response
+        result = manager.handle_elicitation_response(
+            elicitation_id="elicit_123",
+            action="accept",
+            data={"name": "John"}
+        )
+
+        assert result is True
+
+    @pytest.mark.asyncio
+    async def test_handle_elicitation_response_decline(self):
+        """Test handling a decline response."""
+        manager = ElicitationManager()
+        manager.create_elicitation_request(
+            elicitation_id="elicit_123",
+            tool_call_id="tool_456",
+            tool_name="test_tool",
+            message="Enter your name",
+            response_schema={"type": "object"}
+        )
+
+        # Handle decline response
+        result = manager.handle_elicitation_response(
+            elicitation_id="elicit_123",
+            action="decline",
+            data=None
+        )
+
+        assert result is True
+
+    @pytest.mark.asyncio
+    async def test_handle_elicitation_response_cancel(self):
+        """Test handling a cancel response."""
+        manager = ElicitationManager()
+        manager.create_elicitation_request(
+            elicitation_id="elicit_123",
+            tool_call_id="tool_456",
+            tool_name="test_tool",
+            message="Enter your name",
+            response_schema={"type": "object"}
+        )
+
+        # Handle cancel response
+        result = manager.handle_elicitation_response(
+            elicitation_id="elicit_123",
+            action="cancel",
+            data=None
+        )
+
+        assert result is True
+
+    def test_handle_unknown_elicitation(self):
+        """Test handling response for unknown elicitation."""
+        manager = ElicitationManager()
+
+        # Try to handle response for non-existent elicitation
+        result = manager.handle_elicitation_response(
+            elicitation_id="unknown_123",
+            action="accept",
+            data={"name": "John"}
+        )
+
+        assert result is False
+
+    @pytest.mark.asyncio
+    async def test_cleanup_request(self):
+        """Test cleaning up an elicitation request."""
+        manager = ElicitationManager()
+        manager.create_elicitation_request(
+            elicitation_id="elicit_123",
+            tool_call_id="tool_456",
+            tool_name="test_tool",
+            message="Enter your name",
+            response_schema={"type": "object"}
+        )
+
+        # Verify request exists
+        assert "elicit_123" in manager.get_all_pending_requests()
+
+        # Cleanup the request
+        manager.cleanup_request("elicit_123")
+
+        # Verify request is removed
+        assert "elicit_123" not in manager.get_all_pending_requests()
+
+    @pytest.mark.asyncio
+    async def test_get_pending_request(self):
+        """Test retrieving a pending request."""
+        manager = ElicitationManager()
+        manager.create_elicitation_request(
+            elicitation_id="elicit_123",
+            tool_call_id="tool_456",
+            tool_name="test_tool",
+            message="Enter your name",
+            response_schema={"type": "object"}
+        )
+
+        # Retrieve the request
+        retrieved = manager.get_pending_request("elicit_123")
+
+        assert retrieved is not None
+        assert retrieved.elicitation_id == "elicit_123"
+        assert retrieved.tool_call_id == "tool_456"
+
+    @pytest.mark.asyncio
+    async def test_get_pending_request_not_found(self):
+        """Test retrieving non-existent request."""
+        manager = ElicitationManager()
+
+        # Try to get non-existent request
+        retrieved = manager.get_pending_request("unknown_123")
+
+        assert retrieved is None
+
+    @pytest.mark.asyncio
+    async def test_cancel_all_requests(self):
+        """Test cancelling all pending requests."""
+        manager = ElicitationManager()
+
+        # Create multiple requests
+        manager.create_elicitation_request(
+            elicitation_id="elicit_1",
+            tool_call_id="tool_1",
+            tool_name="test_tool",
+            message="Request 1",
+            response_schema={"type": "object"}
+        )
+        manager.create_elicitation_request(
+            elicitation_id="elicit_2",
+            tool_call_id="tool_2",
+            tool_name="test_tool",
+            message="Request 2",
+            response_schema={"type": "object"}
+        )
+
+        # Verify both exist
+        assert len(manager.get_all_pending_requests()) == 2
+
+        # Cancel all requests
+        manager.cancel_all_requests()
+
+        # Verify all are removed
+        assert len(manager.get_all_pending_requests()) == 0
+
+    def test_get_elicitation_manager_singleton(self):
+        """Test that get_elicitation_manager returns singleton."""
+        manager1 = get_elicitation_manager()
+        manager2 = get_elicitation_manager()
+
+        assert manager1 is manager2
+
+
+class TestElicitationManagerIntegration:
+    """Integration tests for ElicitationManager."""
+
+    @pytest.mark.asyncio
+    async def test_full_elicitation_flow(self):
+        """Test complete elicitation flow from request to response."""
+        manager = ElicitationManager()
+
+        # Create request
+        request = manager.create_elicitation_request(
+            elicitation_id="elicit_123",
+            tool_call_id="tool_456",
+            tool_name="test_tool",
+            message="Enter your information",
+            response_schema={
+                "type": "object",
+                "properties": {
+                    "name": {"type": "string"},
+                    "age": {"type": "integer"}
+                }
+            }
+        )
+
+        # Simulate async handling
+        async def simulate_user_response():
+            # Wait a bit to simulate user thinking
+            await asyncio.sleep(0.1)
+            # User responds
+            manager.handle_elicitation_response(
+                elicitation_id="elicit_123",
+                action="accept",
+                data={"name": "Alice", "age": 30}
+            )
+
+        # Start the simulation
+        asyncio.create_task(simulate_user_response())
+
+        # Wait for response
+        response = await request.wait_for_response(timeout=2.0)
+
+        assert response["action"] == "accept"
+        assert response["data"]["name"] == "Alice"
+        assert response["data"]["age"] == 30
+
+        # Cleanup
+        manager.cleanup_request("elicit_123")
+        assert "elicit_123" not in manager.get_all_pending_requests()
+
+    @pytest.mark.asyncio
+    async def test_multi_turn_elicitation(self):
+        """Test handling multiple sequential elicitation requests."""
+        manager = ElicitationManager()
+
+        # First elicitation
+        request1 = manager.create_elicitation_request(
+            elicitation_id="elicit_1",
+            tool_call_id="tool_456",
+            tool_name="test_tool",
+            message="Enter your name",
+            response_schema={"type": "object"}
+        )
+
+        # Immediately respond to first
+        manager.handle_elicitation_response(
+            elicitation_id="elicit_1",
+            action="accept",
+            data={"name": "Bob"}
+        )
+
+        response1 = await request1.wait_for_response(timeout=1.0)
+        assert response1["action"] == "accept"
+
+        manager.cleanup_request("elicit_1")
+
+        # Second elicitation
+        request2 = manager.create_elicitation_request(
+            elicitation_id="elicit_2",
+            tool_call_id="tool_456",
+            tool_name="test_tool",
+            message="Enter your age",
+            response_schema={"type": "object"}
+        )
+
+        # Respond to second
+        manager.handle_elicitation_response(
+            elicitation_id="elicit_2",
+            action="accept",
+            data={"age": 25}
+        )
+
+        response2 = await request2.wait_for_response(timeout=1.0)
+        assert response2["action"] == "accept"
+        assert response2["data"]["age"] == 25
+
+        manager.cleanup_request("elicit_2")
+
+    @pytest.mark.asyncio
+    async def test_elicitation_with_decline(self):
+        """Test elicitation flow when user declines."""
+        manager = ElicitationManager()
+
+        request = manager.create_elicitation_request(
+            elicitation_id="elicit_123",
+            tool_call_id="tool_456",
+            tool_name="test_tool",
+            message="Enter optional information",
+            response_schema={"type": "object"}
+        )
+
+        # User declines
+        manager.handle_elicitation_response(
+            elicitation_id="elicit_123",
+            action="decline",
+            data=None
+        )
+
+        response = await request.wait_for_response(timeout=1.0)
+
+        assert response["action"] == "decline"
+        assert response["data"] is None
+
+        manager.cleanup_request("elicit_123")

atlas/tests/test_elicitation_routing.py

@@ -0,0 +1,296 @@
+"""
+Tests for MCP elicitation routing functionality.
+
+Tests the dictionary-based routing system that allows elicitation requests
+from MCP tools to reach the correct WebSocket connection across async tasks.
+"""
+
+from unittest.mock import AsyncMock, Mock, patch
+
+import pytest
+
+from atlas.domain.messages.models import ToolCall
+
+
+class TestElicitationRouting:
+    """Test elicitation routing context management."""
+
+    @pytest.fixture
+    def mock_tool_call(self):
+        """Create a mock ToolCall object."""
+        return ToolCall(
+            id="test_call_123",
+            name="elicitation_demo_get_user_name",
+            arguments={}
+        )
+
+    @pytest.fixture
+    def mock_update_callback(self):
+        """Create a mock update callback."""
+        return AsyncMock()
+
+    @pytest.fixture
+    def manager(self):
+        """Create a MCPToolManager instance for testing."""
+        from atlas.modules.mcp_tools.client import MCPToolManager
+        return MCPToolManager(config_path="/tmp/nonexistent_mcp_test.json")
+
+    @pytest.mark.asyncio
+    async def test_elicitation_context_sets_routing(self, manager, mock_tool_call, mock_update_callback):
+        """Test that elicitation context correctly sets routing in dictionary."""
+        from atlas.modules.mcp_tools.client import _ELICITATION_ROUTING
+
+        # Clear routing before test
+        _ELICITATION_ROUTING.clear()
+
+        server_name = "test_server"
+        routing_key = (server_name, mock_tool_call.id)
+
+        # Use the context manager
+        async with manager._use_elicitation_context(server_name, mock_tool_call, mock_update_callback):
+            # Inside context: routing should exist with composite key
+            assert routing_key in _ELICITATION_ROUTING
+            routing = _ELICITATION_ROUTING[routing_key]
+            assert routing.server_name == server_name
+            assert routing.tool_call == mock_tool_call
+            assert routing.update_cb == mock_update_callback
+
+        # After context: routing should be cleaned up
+        assert routing_key not in _ELICITATION_ROUTING
+
+    @pytest.mark.asyncio
+    async def test_elicitation_routing_cleanup_on_error(self, manager, mock_tool_call, mock_update_callback):
+        """Test that routing is cleaned up even if error occurs."""
+        from atlas.modules.mcp_tools.client import _ELICITATION_ROUTING
+
+        _ELICITATION_ROUTING.clear()
+
+        server_name = "test_server"
+        routing_key = (server_name, mock_tool_call.id)
+
+        # Simulate an error inside the context
+        with pytest.raises(RuntimeError):
+            async with manager._use_elicitation_context(server_name, mock_tool_call, mock_update_callback):
+                assert routing_key in _ELICITATION_ROUTING
+                raise RuntimeError("Simulated error")
+
+        # Routing should still be cleaned up
+        assert routing_key not in _ELICITATION_ROUTING
+
+    @pytest.mark.asyncio
+    async def test_multiple_servers_routing(self, manager, mock_update_callback):
+        """Test that multiple servers can have separate routing contexts."""
+        from atlas.modules.mcp_tools.client import _ELICITATION_ROUTING
+
+        _ELICITATION_ROUTING.clear()
+
+        tool_call_1 = ToolCall(id="call_1", name="tool_1", arguments={})
+        tool_call_2 = ToolCall(id="call_2", name="tool_2", arguments={})
+        routing_key_1 = ("server_1", "call_1")
+        routing_key_2 = ("server_2", "call_2")
+
+        # Create contexts for two different servers
+        async with manager._use_elicitation_context("server_1", tool_call_1, mock_update_callback):
+            async with manager._use_elicitation_context("server_2", tool_call_2, mock_update_callback):
+                # Both should exist simultaneously
+                assert routing_key_1 in _ELICITATION_ROUTING
+                assert routing_key_2 in _ELICITATION_ROUTING
+                assert _ELICITATION_ROUTING[routing_key_1].tool_call == tool_call_1
+                assert _ELICITATION_ROUTING[routing_key_2].tool_call == tool_call_2
+
+            # server_2 cleaned up, server_1 still exists
+            assert routing_key_1 in _ELICITATION_ROUTING
+            assert routing_key_2 not in _ELICITATION_ROUTING
+
+        # Both cleaned up
+        assert routing_key_1 not in _ELICITATION_ROUTING
+        assert routing_key_2 not in _ELICITATION_ROUTING
+
+    @pytest.mark.asyncio
+    async def test_elicitation_with_none_callback(self, manager, mock_tool_call):
+        """Test elicitation context with None callback (should still work)."""
+        from atlas.modules.mcp_tools.client import _ELICITATION_ROUTING
+
+        _ELICITATION_ROUTING.clear()
+
+        server_name = "test_server"
+        routing_key = (server_name, mock_tool_call.id)
+
+        # Use context with None callback
+        async with manager._use_elicitation_context(server_name, mock_tool_call, None):
+            routing = _ELICITATION_ROUTING[routing_key]
+            assert routing.update_cb is None
+
+        assert routing_key not in _ELICITATION_ROUTING
+
+
+class TestElicitationHandler:
+    """Test per-server elicitation handler creation."""
+
+    @pytest.fixture
+    def manager(self):
+        """Create a MCPToolManager instance for testing."""
+        from atlas.modules.mcp_tools.client import MCPToolManager
+        return MCPToolManager(config_path="/tmp/nonexistent_mcp_test.json")
+
+    @pytest.mark.asyncio
+    async def test_handler_creation_captures_server_name(self, manager):
+        """Test that handler closure captures the correct server_name."""
+
+        # Create handlers for different servers
+        handler_1 = manager._create_elicitation_handler("server_1")
+        handler_2 = manager._create_elicitation_handler("server_2")
+
+        # Handlers should be different functions (different closures)
+        assert handler_1 != handler_2
+        assert callable(handler_1)
+        assert callable(handler_2)
+
+    @pytest.mark.asyncio
+    async def test_handler_returns_cancel_when_no_routing(self, manager):
+        """Test that handler returns cancel when routing not found."""
+        from fastmcp.client.elicitation import ElicitResult
+
+        from atlas.modules.mcp_tools.client import _ELICITATION_ROUTING
+
+        _ELICITATION_ROUTING.clear()
+
+        handler = manager._create_elicitation_handler("test_server")
+
+        # Call handler with no routing set
+        result = await handler("Test message", str, None, None)
+
+        assert isinstance(result, ElicitResult)
+        assert result.action == "cancel"
+        assert result.content is None
+
+    @pytest.mark.asyncio
+    async def test_handler_returns_cancel_when_no_update_cb(self, manager):
+        """Test that handler returns cancel when update_cb is None."""
+        from fastmcp.client.elicitation import ElicitResult
+
+        from atlas.domain.messages.models import ToolCall
+        from atlas.modules.mcp_tools.client import _ELICITATION_ROUTING, _ElicitationRoutingContext
+
+        _ELICITATION_ROUTING.clear()
+
+        server_name = "test_server"
+        tool_call = ToolCall(id="call_123", name="test_tool", arguments={})
+        routing_key = (server_name, tool_call.id)
+
+        # Set routing with None callback using composite key
+        _ELICITATION_ROUTING[routing_key] = _ElicitationRoutingContext(
+            server_name=server_name,
+            tool_call=tool_call,
+            update_cb=None
+        )
+
+        handler = manager._create_elicitation_handler(server_name)
+        result = await handler("Test message", str, None, None)
+
+        assert isinstance(result, ElicitResult)
+        assert result.action == "cancel"
+        assert result.content is None
+
+
+class TestElicitationIntegration:
+    """Integration tests for elicitation workflow."""
+
+    @pytest.fixture
+    def manager(self):
+        """Create a MCPToolManager instance for testing."""
+        from atlas.modules.mcp_tools.client import MCPToolManager
+        return MCPToolManager(config_path="/tmp/nonexistent_mcp_test.json")
+
+    @pytest.mark.asyncio
+    async def test_elicitation_request_sent_to_callback(self, manager):
+        """Test that elicitation request is sent to update callback."""
+        from atlas.domain.messages.models import ToolCall
+        from atlas.modules.mcp_tools.client import _ELICITATION_ROUTING, _ElicitationRoutingContext
+
+        _ELICITATION_ROUTING.clear()
+
+        server_name = "test_server"
+        tool_call = ToolCall(id="call_123", name="test_tool", arguments={})
+        mock_callback = AsyncMock()
+        routing_key = (server_name, tool_call.id)
+
+        # Set routing with mock callback using composite key
+        _ELICITATION_ROUTING[routing_key] = _ElicitationRoutingContext(
+            server_name=server_name,
+            tool_call=tool_call,
+            update_cb=mock_callback
+        )
+
+        handler = manager._create_elicitation_handler(server_name)
+
+        # Mock elicitation manager
+        with patch('atlas.application.chat.elicitation_manager.get_elicitation_manager') as mock_get_mgr:
+            mock_elicit_mgr = Mock()
+            mock_request = AsyncMock()
+            mock_request.wait_for_response = AsyncMock(return_value={
+                "action": "accept",
+                "data": "test_value"
+            })
+            mock_elicit_mgr.create_elicitation_request = Mock(return_value=mock_request)
+            mock_elicit_mgr.cleanup_request = Mock()
+            mock_get_mgr.return_value = mock_elicit_mgr
+
+            result = await handler("What's your name?", str, None, None)
+
+            # Verify callback was called with elicitation_request
+            mock_callback.assert_called_once()
+            call_args = mock_callback.call_args[0][0]
+            assert call_args["type"] == "elicitation_request"
+            assert call_args["message"] == "What's your name?"
+            assert call_args["tool_call_id"] == "call_123"
+
+            # Verify result
+            assert result.action == "accept"
+            assert result.content == {"value": "test_value"}
+
+    @pytest.mark.asyncio
+    async def test_elicitation_accept_no_data_returns_empty_object(self, manager):
+        """Test approval-only elicitation returns empty object on accept.
+
+        FastMCP validation for response_type=None expects an empty response object.
+        Some UIs send placeholder payloads like {'none': ''}; we must not forward them.
+        """
+        from atlas.domain.messages.models import ToolCall
+        from atlas.modules.mcp_tools.client import _ELICITATION_ROUTING, _ElicitationRoutingContext
+
+        _ELICITATION_ROUTING.clear()
+
+        server_name = "test_server"
+        tool_call = ToolCall(id="call_123", name="test_tool", arguments={})
+        mock_callback = AsyncMock()
+        routing_key = (server_name, tool_call.id)
+
+        _ELICITATION_ROUTING[routing_key] = _ElicitationRoutingContext(
+            server_name=server_name,
+            tool_call=tool_call,
+            update_cb=mock_callback,
+        )
+
+        handler = manager._create_elicitation_handler(server_name)
+
+        with patch('atlas.application.chat.elicitation_manager.get_elicitation_manager') as mock_get_mgr:
+            mock_elicit_mgr = Mock()
+            mock_request = AsyncMock()
+            mock_request.wait_for_response = AsyncMock(return_value={
+                "action": "accept",
+                "data": {"none": ""},
+            })
+            mock_elicit_mgr.create_elicitation_request = Mock(return_value=mock_request)
+            mock_elicit_mgr.cleanup_request = Mock()
+            mock_get_mgr.return_value = mock_elicit_mgr
+
+            result = await handler(
+                "Are you sure you want to delete this item?",
+                None,
+                None,
+                None,
+            )
+
+            assert result.action == "accept"
+            assert result.content == {}