atlas-chat 0.1.0__py3-none-any.whl

atlas/tests/test_mcp_image_content.py

@@ -0,0 +1,424 @@
+"""Tests for ImageContent handling in MCP tool results.
+
+These tests verify that Atlas can extract and process ImageContent items
+from MCP tool responses and convert them to artifacts for display.
+"""
+
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from atlas.domain.messages.models import ToolCall
+from atlas.modules.mcp_tools.client import MCPToolManager
+
+
+class MockImageContent:
+    """Mock for MCP ImageContent item."""
+    def __init__(self, data: str, mime_type: str = "image/png"):
+        self.type = "image"
+        self.data = data
+        self.mimeType = mime_type
+
+
+class MockTextContent:
+    """Mock for MCP text content item."""
+    def __init__(self, text: str):
+        self.type = "text"
+        self.text = text
+
+
+class MockMCPResultWithImage:
+    """Mock MCP result that includes ImageContent in content array."""
+    def __init__(self, image_data: str, mime_type: str = "image/png"):
+        self.content = [MockImageContent(image_data, mime_type)]
+        self.structured_content = None
+        self.data = None
+        self.is_error = False
+
+
+class MockMCPResultWithMultipleImages:
+    """Mock MCP result with multiple ImageContent items."""
+    def __init__(self, images: list):
+        self.content = [MockImageContent(img["data"], img["mime"]) for img in images]
+        self.structured_content = None
+        self.data = None
+        self.is_error = False
+
+
+class MockMCPResultWithMixedContent:
+    """Mock MCP result with both TextContent and ImageContent."""
+    def __init__(self, text: str, image_data: str):
+        self.content = [
+            MockTextContent(text),
+            MockImageContent(image_data)
+        ]
+        self.structured_content = None
+        self.data = None
+        self.is_error = False
+
+
+class TestImageContentHandling:
+    """Tests for extracting ImageContent from MCP tool results."""
+
+    @pytest.mark.asyncio
+    async def test_extract_single_image_content(self):
+        """Test extraction of a single ImageContent item."""
+        manager = MCPToolManager.__new__(MCPToolManager)
+
+        # Mock tool object
+        class MockTool:
+            def __init__(self, name):
+                self.name = name
+
+        # Create a tool call
+        tool_call = ToolCall(
+            id="test-call-1",
+            name="generate_image",
+            arguments={}
+        )
+
+        # Mock the call_tool to return ImageContent
+        image_b64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="
+        raw_result = MockMCPResultWithImage(image_b64, "image/png")
+
+        with patch.object(manager, 'call_tool', new_callable=AsyncMock) as mock_call:
+            mock_call.return_value = raw_result
+
+            # Mock _tool_index with correct structure
+            manager._tool_index = {
+                "generate_image": {
+                    "server": "test-server",
+                    "tool": MockTool("generate_image")
+                }
+            }
+
+            result = await manager.execute_tool(tool_call, context={})
+
+            # Verify artifacts were created
+            assert result.artifacts is not None
+            assert len(result.artifacts) == 1
+
+            artifact = result.artifacts[0]
+            assert artifact["name"] == "mcp_image_0.png"
+            assert artifact["b64"] == image_b64
+            assert artifact["mime"] == "image/png"
+            assert artifact["viewer"] == "image"
+            assert "generate_image" in artifact["description"]
+
+            # Verify display config was auto-created
+            assert result.display_config is not None
+            assert result.display_config["primary_file"] == "mcp_image_0.png"
+            assert result.display_config["open_canvas"] is True
+
+    @pytest.mark.asyncio
+    async def test_extract_multiple_image_contents(self):
+        """Test extraction of multiple ImageContent items."""
+        manager = MCPToolManager.__new__(MCPToolManager)
+
+        # Mock tool object
+        class MockTool:
+            def __init__(self, name):
+                self.name = name
+
+        tool_call = ToolCall(
+            id="test-call-2",
+            name="generate_multiple",
+            arguments={}
+        )
+
+        # Use valid base64 encoded strings
+        images = [
+            {"data": "aW1hZ2UgZGF0YSAxCg==", "mime": "image/png"},
+            {"data": "aW1hZ2UgZGF0YSAyCg==", "mime": "image/jpeg"},
+            {"data": "aW1hZ2UgZGF0YSAzCg==", "mime": "image/gif"}
+        ]
+        raw_result = MockMCPResultWithMultipleImages(images)
+
+        with patch.object(manager, 'call_tool', new_callable=AsyncMock) as mock_call:
+            mock_call.return_value = raw_result
+            manager._tool_index = {
+                "generate_multiple": {
+                    "server": "test-server",
+                    "tool": MockTool("generate_multiple")
+                }
+            }
+
+            result = await manager.execute_tool(tool_call, context={})
+
+            # Verify all images were extracted
+            assert len(result.artifacts) == 3
+
+            # Check each artifact
+            for i, (artifact, img) in enumerate(zip(result.artifacts, images)):
+                expected_ext = img["mime"].split("/")[-1]
+                assert artifact["name"] == f"mcp_image_{i}.{expected_ext}"
+                assert artifact["b64"] == img["data"]
+                assert artifact["mime"] == img["mime"]
+                assert artifact["viewer"] == "image"
+
+    @pytest.mark.asyncio
+    async def test_extract_mixed_content(self):
+        """Test extraction when both TextContent and ImageContent are present."""
+        manager = MCPToolManager.__new__(MCPToolManager)
+
+        # Mock tool object
+        class MockTool:
+            def __init__(self, name):
+                self.name = name
+
+        tool_call = ToolCall(
+            id="test-call-3",
+            name="mixed_tool",
+            arguments={}
+        )
+
+        text = "Here is your visualization"
+        # Use valid base64 encoded string
+        image_b64 = "aW1hZ2VkYXRhCg=="
+        raw_result = MockMCPResultWithMixedContent(text, image_b64)
+
+        with patch.object(manager, 'call_tool', new_callable=AsyncMock) as mock_call:
+            mock_call.return_value = raw_result
+            manager._tool_index = {
+                "mixed_tool": {
+                    "server": "test-server",
+                    "tool": MockTool("mixed_tool")
+                }
+            }
+
+            result = await manager.execute_tool(tool_call, context={})
+
+            # Verify image was extracted (uses image counter, so first image is image_0)
+            assert len(result.artifacts) == 1
+            artifact = result.artifacts[0]
+            assert artifact["name"] == "mcp_image_0.png"
+            assert artifact["b64"] == image_b64
+
+            # Verify the text content was extracted and included in result.content
+            # The content should be JSON containing the text in "results"
+            import json
+            content_dict = json.loads(result.content)
+            assert "results" in content_dict
+            assert text in content_dict["results"]
+
+    @pytest.mark.asyncio
+    async def test_no_image_content(self):
+        """Test that non-image content doesn't create artifacts."""
+        manager = MCPToolManager.__new__(MCPToolManager)
+
+        # Mock tool object
+        class MockTool:
+            def __init__(self, name):
+                self.name = name
+
+        tool_call = ToolCall(
+            id="test-call-4",
+            name="text_only",
+            arguments={}
+        )
+
+        # Create a result with only text content
+        class MockTextOnlyResult:
+            def __init__(self):
+                self.content = [MockTextContent("Just text")]
+                self.structured_content = None
+                self.data = None
+                self.is_error = False
+
+        raw_result = MockTextOnlyResult()
+
+        with patch.object(manager, 'call_tool', new_callable=AsyncMock) as mock_call:
+            mock_call.return_value = raw_result
+            manager._tool_index = {
+                "text_only": {
+                    "server": "test-server",
+                    "tool": MockTool("text_only")
+                }
+            }
+
+            result = await manager.execute_tool(tool_call, context={})
+
+            # Verify no artifacts were created
+            assert len(result.artifacts) == 0
+            # Display config should not be auto-created
+            assert result.display_config is None
+
+    @pytest.mark.asyncio
+    async def test_image_content_missing_data(self):
+        """Test that ImageContent with None/missing data is skipped."""
+        manager = MCPToolManager.__new__(MCPToolManager)
+
+        class MockTool:
+            def __init__(self, name):
+                self.name = name
+
+        tool_call = ToolCall(
+            id="test-call-5",
+            name="missing_data",
+            arguments={}
+        )
+
+        # Create ImageContent with missing data
+        class MockImageContentNoData:
+            def __init__(self):
+                self.type = "image"
+                self.data = None
+                self.mimeType = "image/png"
+
+        class MockResult:
+            def __init__(self):
+                self.content = [MockImageContentNoData()]
+                self.structured_content = None
+                self.data = None
+                self.is_error = False
+
+        raw_result = MockResult()
+
+        with patch.object(manager, 'call_tool', new_callable=AsyncMock) as mock_call:
+            mock_call.return_value = raw_result
+            manager._tool_index = {
+                "missing_data": {
+                    "server": "test-server",
+                    "tool": MockTool("missing_data")
+                }
+            }
+
+            result = await manager.execute_tool(tool_call, context={})
+
+            # No artifacts should be created when data is missing
+            assert len(result.artifacts) == 0
+
+    @pytest.mark.asyncio
+    async def test_image_content_missing_mime_type(self):
+        """Test that ImageContent with None/missing mimeType is skipped."""
+        manager = MCPToolManager.__new__(MCPToolManager)
+
+        class MockTool:
+            def __init__(self, name):
+                self.name = name
+
+        tool_call = ToolCall(
+            id="test-call-6",
+            name="missing_mime",
+            arguments={}
+        )
+
+        # Create ImageContent with missing mime type
+        class MockImageContentNoMime:
+            def __init__(self):
+                self.type = "image"
+                self.data = "SGVsbG8gV29ybGQ="  # Valid base64
+                self.mimeType = None
+
+        class MockResult:
+            def __init__(self):
+                self.content = [MockImageContentNoMime()]
+                self.structured_content = None
+                self.data = None
+                self.is_error = False
+
+        raw_result = MockResult()
+
+        with patch.object(manager, 'call_tool', new_callable=AsyncMock) as mock_call:
+            mock_call.return_value = raw_result
+            manager._tool_index = {
+                "missing_mime": {
+                    "server": "test-server",
+                    "tool": MockTool("missing_mime")
+                }
+            }
+
+            result = await manager.execute_tool(tool_call, context={})
+
+            # No artifacts should be created when mimeType is missing
+            assert len(result.artifacts) == 0
+
+    @pytest.mark.asyncio
+    async def test_image_content_invalid_mime_type(self):
+        """Test that ImageContent with unsupported mime type is skipped."""
+        manager = MCPToolManager.__new__(MCPToolManager)
+
+        class MockTool:
+            def __init__(self, name):
+                self.name = name
+
+        tool_call = ToolCall(
+            id="test-call-7",
+            name="bad_mime",
+            arguments={}
+        )
+
+        # Create ImageContent with unsupported mime type
+        class MockImageContentBadMime:
+            def __init__(self):
+                self.type = "image"
+                self.data = "SGVsbG8gV29ybGQ="  # Valid base64
+                self.mimeType = "application/octet-stream"
+
+        class MockResult:
+            def __init__(self):
+                self.content = [MockImageContentBadMime()]
+                self.structured_content = None
+                self.data = None
+                self.is_error = False
+
+        raw_result = MockResult()
+
+        with patch.object(manager, 'call_tool', new_callable=AsyncMock) as mock_call:
+            mock_call.return_value = raw_result
+            manager._tool_index = {
+                "bad_mime": {
+                    "server": "test-server",
+                    "tool": MockTool("bad_mime")
+                }
+            }
+
+            result = await manager.execute_tool(tool_call, context={})
+
+            # No artifacts should be created for unsupported mime type
+            assert len(result.artifacts) == 0
+
+    @pytest.mark.asyncio
+    async def test_image_content_invalid_base64(self):
+        """Test that ImageContent with invalid base64 data is skipped."""
+        manager = MCPToolManager.__new__(MCPToolManager)
+
+        class MockTool:
+            def __init__(self, name):
+                self.name = name
+
+        tool_call = ToolCall(
+            id="test-call-8",
+            name="bad_base64",
+            arguments={}
+        )
+
+        # Create ImageContent with invalid base64
+        class MockImageContentBadB64:
+            def __init__(self):
+                self.type = "image"
+                self.data = "not-valid-base64!!!"
+                self.mimeType = "image/png"
+
+        class MockResult:
+            def __init__(self):
+                self.content = [MockImageContentBadB64()]
+                self.structured_content = None
+                self.data = None
+                self.is_error = False
+
+        raw_result = MockResult()
+
+        with patch.object(manager, 'call_tool', new_callable=AsyncMock) as mock_call:
+            mock_call.return_value = raw_result
+            manager._tool_index = {
+                "bad_base64": {
+                    "server": "test-server",
+                    "tool": MockTool("bad_base64")
+                }
+            }
+
+            result = await manager.execute_tool(tool_call, context={})
+
+            # No artifacts should be created for invalid base64
+            assert len(result.artifacts) == 0

atlas/tests/test_mcp_logging.py

@@ -0,0 +1,172 @@
+"""Tests for MCP server logging functionality.
+
+These tests verify that:
+1. Log handlers are properly created and attached to MCP clients
+2. Log messages are filtered based on configured LOG_LEVEL
+3. Log messages are forwarded to the UI callback when provided
+4. Backend logger receives MCP server logs
+"""
+
+import asyncio
+import logging
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from atlas.modules.mcp_tools.client import MCP_TO_PYTHON_LOG_LEVEL, MCPToolManager
+
+
+class MockLogMessage:
+    """Mock LogMessage from fastmcp.client.logging."""
+    def __init__(self, level: str, msg: str, extra: dict = None):
+        self.level = level
+        self.data = {
+            'msg': msg,
+            'extra': extra or {}
+        }
+
+
+@pytest.mark.asyncio
+class TestMCPLogging:
+    """Tests for MCP logging functionality."""
+
+    async def test_log_level_mapping(self):
+        """Test that MCP log levels are mapped correctly to Python logging levels."""
+        assert MCP_TO_PYTHON_LOG_LEVEL['debug'] == logging.DEBUG
+        assert MCP_TO_PYTHON_LOG_LEVEL['info'] == logging.INFO
+        assert MCP_TO_PYTHON_LOG_LEVEL['notice'] == logging.INFO
+        assert MCP_TO_PYTHON_LOG_LEVEL['warning'] == logging.WARNING
+        assert MCP_TO_PYTHON_LOG_LEVEL['warn'] == logging.WARNING
+        assert MCP_TO_PYTHON_LOG_LEVEL['error'] == logging.ERROR
+        assert MCP_TO_PYTHON_LOG_LEVEL['alert'] == logging.CRITICAL
+        assert MCP_TO_PYTHON_LOG_LEVEL['critical'] == logging.CRITICAL
+        assert MCP_TO_PYTHON_LOG_LEVEL['emergency'] == logging.CRITICAL
+
+    async def test_log_handler_forwards_to_callback(self):
+        """Test that log handler forwards messages to UI callback."""
+        # Create a mock callback
+        mock_callback = AsyncMock()
+
+        with patch.dict('os.environ', {'LOG_LEVEL': 'DEBUG'}):
+            manager = MCPToolManager(log_callback=mock_callback)
+            log_handler = manager._create_log_handler("test_server")
+
+            # Send a log message
+            msg = MockLogMessage('info', 'Test message', {'key': 'value'})
+            await log_handler(msg)
+
+            # Callback should be called with correct parameters
+            mock_callback.assert_called_once()
+            call_args = mock_callback.call_args[0]
+            assert call_args[0] == "test_server"  # server_name
+            assert call_args[1] == "info"  # level
+            assert call_args[2] == "Test message"  # message
+            assert call_args[3] == {'key': 'value'}  # extra
+
+    async def test_log_handler_filters_by_level(self):
+        """Test that log handler respects min_log_level filtering."""
+        mock_callback = AsyncMock()
+
+        # Set minimum level to WARNING
+        with patch.dict('os.environ', {'LOG_LEVEL': 'WARNING'}):
+            manager = MCPToolManager(log_callback=mock_callback)
+            manager._min_log_level = logging.WARNING  # Ensure it's set
+            log_handler = manager._create_log_handler("test_server")
+
+            # Send a DEBUG log (should be filtered out)
+            debug_msg = MockLogMessage('debug', 'Debug message')
+            await log_handler(debug_msg)
+
+            # Callback should NOT be called for DEBUG when level is WARNING
+            mock_callback.assert_not_called()
+
+            # Send an INFO log (should also be filtered out)
+            info_msg = MockLogMessage('info', 'Info message')
+            await log_handler(info_msg)
+
+            # Still should not be called
+            mock_callback.assert_not_called()
+
+            # Send a WARNING log (should pass through)
+            warn_msg = MockLogMessage('warning', 'Warning message')
+            await log_handler(warn_msg)
+
+            # Now callback should be called
+            mock_callback.assert_called_once()
+
+    async def test_set_log_callback(self):
+        """Test that log callback can be set after initialization."""
+        manager = MCPToolManager()
+
+        # Initially no callback
+        assert manager._default_log_callback is None
+
+        # Set a callback
+        mock_callback = AsyncMock()
+        manager.set_log_callback(mock_callback)
+
+        assert manager._default_log_callback is mock_callback
+
+        # Test that it's used
+        log_handler = manager._create_log_handler("test_server")
+
+        msg = MockLogMessage('info', 'Test message')
+        await log_handler(msg)
+
+        mock_callback.assert_called_once()
+
+    async def test_log_handler_handles_callback_errors_gracefully(self):
+        """Test that log handler doesn't crash if callback raises an exception."""
+        # Create a callback that raises an exception
+        mock_callback = AsyncMock(side_effect=Exception("Callback error"))
+
+        manager = MCPToolManager(log_callback=mock_callback)
+        log_handler = manager._create_log_handler("test_server")
+
+        # Send a log message - should not raise despite callback error
+        msg = MockLogMessage('info', 'Test message')
+        # This should not raise an exception
+        await log_handler(msg)
+
+        # Verify the callback was attempted
+        mock_callback.assert_called_once()
+
+    async def test_request_scoped_callback_overrides_default(self):
+        """Request-scoped callback should override the default callback.
+
+        This is the core mechanism preventing cross-user log leakage when MCPToolManager
+        is shared across multiple websocket connections.
+        """
+        default_cb = AsyncMock()
+        request_cb = AsyncMock()
+
+        manager = MCPToolManager(log_callback=default_cb)
+        log_handler = manager._create_log_handler("test_server")
+
+        msg = MockLogMessage('info', 'Scoped message')
+
+        async with manager._use_log_callback(request_cb):
+            await log_handler(msg)
+
+        request_cb.assert_called_once()
+        default_cb.assert_not_called()
+
+    async def test_request_scoped_callbacks_are_isolated_across_tasks(self):
+        """Two concurrent tasks should not receive each other's MCP logs."""
+        cb_a = AsyncMock()
+        cb_b = AsyncMock()
+        manager = MCPToolManager()
+        log_handler = manager._create_log_handler("test_server")
+
+        async def run_with(cb, text):
+            async with manager._use_log_callback(cb):
+                await log_handler(MockLogMessage('info', text))
+
+        await asyncio.gather(
+            run_with(cb_a, "message-a"),
+            run_with(cb_b, "message-b"),
+        )
+
+        cb_a.assert_called_once()
+        cb_b.assert_called_once()
+