atlas-chat 0.1.0__py3-none-any.whl

atlas/tests/test_mcp_client_auth.py

@@ -0,0 +1,331 @@
+"""Unit tests for MCP client authentication methods.
+
+Tests the per-user authentication functionality in MCPToolManager:
+- _requires_user_auth: Check if server requires user authentication
+- _get_user_client: Get or create user-specific client with token
+- Cache validation and invalidation
+
+Updated: 2025-01-23
+"""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+
+class TestRequiresUserAuth:
+    """Test _requires_user_auth method."""
+
+    def test_requires_user_auth_for_jwt(self):
+        """JWT auth_type should require user auth."""
+        from atlas.modules.mcp_tools.client import MCPToolManager
+
+        manager = MCPToolManager.__new__(MCPToolManager)
+        manager.servers_config = {"test-server": {"auth_type": "jwt"}}
+
+        assert manager._requires_user_auth("test-server") is True
+
+    def test_requires_user_auth_for_oauth(self):
+        """OAuth auth_type should require user auth."""
+        from atlas.modules.mcp_tools.client import MCPToolManager
+
+        manager = MCPToolManager.__new__(MCPToolManager)
+        manager.servers_config = {"test-server": {"auth_type": "oauth"}}
+
+        assert manager._requires_user_auth("test-server") is True
+
+    def test_requires_user_auth_for_bearer(self):
+        """Bearer auth_type should require user auth."""
+        from atlas.modules.mcp_tools.client import MCPToolManager
+
+        manager = MCPToolManager.__new__(MCPToolManager)
+        manager.servers_config = {"test-server": {"auth_type": "bearer"}}
+
+        assert manager._requires_user_auth("test-server") is True
+
+    def test_requires_user_auth_for_api_key(self):
+        """API key auth_type should require user auth."""
+        from atlas.modules.mcp_tools.client import MCPToolManager
+
+        manager = MCPToolManager.__new__(MCPToolManager)
+        manager.servers_config = {"test-server": {"auth_type": "api_key"}}
+
+        assert manager._requires_user_auth("test-server") is True
+
+    def test_no_user_auth_for_none(self):
+        """None auth_type should not require user auth."""
+        from atlas.modules.mcp_tools.client import MCPToolManager
+
+        manager = MCPToolManager.__new__(MCPToolManager)
+        manager.servers_config = {"test-server": {"auth_type": "none"}}
+
+        assert manager._requires_user_auth("test-server") is False
+
+    def test_no_user_auth_when_missing(self):
+        """Missing auth_type should not require user auth (defaults to none)."""
+        from atlas.modules.mcp_tools.client import MCPToolManager
+
+        manager = MCPToolManager.__new__(MCPToolManager)
+        manager.servers_config = {"test-server": {}}
+
+        assert manager._requires_user_auth("test-server") is False
+
+    def test_no_user_auth_for_unknown_server(self):
+        """Unknown server should not require user auth."""
+        from atlas.modules.mcp_tools.client import MCPToolManager
+
+        manager = MCPToolManager.__new__(MCPToolManager)
+        manager.servers_config = {}
+
+        assert manager._requires_user_auth("unknown-server") is False
+
+
+class TestGetUserClient:
+    """Test _get_user_client method."""
+
+    @pytest.fixture
+    def manager(self):
+        """Create a mock MCPToolManager for testing."""
+        import asyncio
+
+        from atlas.modules.mcp_tools.client import MCPToolManager
+
+        manager = MCPToolManager.__new__(MCPToolManager)
+        manager.servers_config = {
+            "test-server": {
+                "auth_type": "api_key",
+                "url": "http://localhost:8080"
+            }
+        }
+        manager._user_clients = {}
+        manager._user_clients_lock = asyncio.Lock()
+        manager._create_log_handler = MagicMock(return_value=None)
+        manager._create_elicitation_handler = MagicMock(return_value=None)
+        manager._create_sampling_handler = MagicMock(return_value=None)
+        return manager
+
+    @pytest.mark.asyncio
+    async def test_returns_none_without_token(self, manager):
+        """Should return None when user has no token stored."""
+        with patch("atlas.modules.mcp_tools.token_storage.get_token_storage") as mock_storage:
+            mock_token_storage = MagicMock()
+            mock_token_storage.get_valid_token.return_value = None
+            mock_storage.return_value = mock_token_storage
+
+            result = await manager._get_user_client("test-server", "user@example.com")
+
+            assert result is None
+
+    @pytest.mark.asyncio
+    async def test_creates_client_with_api_key_header(self, manager):
+        """Should create client with custom header for API key auth type."""
+        with patch("atlas.modules.mcp_tools.token_storage.get_token_storage") as mock_storage, \
+             patch("atlas.modules.mcp_tools.client.Client") as mock_client_class, \
+             patch("atlas.modules.mcp_tools.client.StreamableHttpTransport") as mock_transport_class:
+
+            # Mock token storage
+            mock_token = MagicMock()
+            mock_token.token_value = "test-api-key-123"
+            mock_token_storage = MagicMock()
+            mock_token_storage.get_valid_token.return_value = mock_token
+            mock_storage.return_value = mock_token_storage
+
+            # Mock transport and Client constructor
+            mock_transport = MagicMock()
+            mock_transport_class.return_value = mock_transport
+            mock_client = MagicMock()
+            mock_client_class.return_value = mock_client
+
+            result = await manager._get_user_client("test-server", "user@example.com")
+
+            assert result is mock_client
+            # Verify StreamableHttpTransport was created with custom header
+            mock_transport_class.assert_called_once()
+            transport_call_kwargs = mock_transport_class.call_args
+            assert transport_call_kwargs[1]["headers"] == {"X-API-Key": "test-api-key-123"}
+            # Verify Client was created with transport
+            mock_client_class.assert_called_once()
+            client_call_kwargs = mock_client_class.call_args
+            assert client_call_kwargs[1]["transport"] is mock_transport
+
+    @pytest.mark.asyncio
+    async def test_creates_client_with_bearer_token(self):
+        """Should create client with auth parameter for bearer auth type."""
+        import asyncio
+
+        from atlas.modules.mcp_tools.client import MCPToolManager
+
+        manager = MCPToolManager.__new__(MCPToolManager)
+        manager.servers_config = {
+            "bearer-server": {
+                "auth_type": "bearer",
+                "url": "http://localhost:8080"
+            }
+        }
+        manager._user_clients = {}
+        manager._user_clients_lock = asyncio.Lock()
+        manager._create_log_handler = MagicMock(return_value=None)
+        manager._create_elicitation_handler = MagicMock(return_value=None)
+        manager._create_sampling_handler = MagicMock(return_value=None)
+
+        with patch("atlas.modules.mcp_tools.token_storage.get_token_storage") as mock_storage, \
+             patch("atlas.modules.mcp_tools.client.Client") as mock_client_class:
+
+            # Mock token storage
+            mock_token = MagicMock()
+            mock_token.token_value = "bearer-token-123"
+            mock_token_storage = MagicMock()
+            mock_token_storage.get_valid_token.return_value = mock_token
+            mock_storage.return_value = mock_token_storage
+
+            # Mock Client constructor
+            mock_client = MagicMock()
+            mock_client_class.return_value = mock_client
+
+            result = await manager._get_user_client("bearer-server", "user@example.com")
+
+            assert result is mock_client
+            mock_client_class.assert_called_once()
+            # Verify token was passed as auth parameter (not via transport)
+            call_args = mock_client_class.call_args
+            assert call_args[0][0] == "http://localhost:8080"  # URL as first positional arg
+            assert call_args[1]["auth"] == "bearer-token-123"
+
+    @pytest.mark.asyncio
+    async def test_uses_custom_auth_header_name(self):
+        """Should use custom auth_header from config for API key auth."""
+        import asyncio
+
+        from atlas.modules.mcp_tools.client import MCPToolManager
+
+        manager = MCPToolManager.__new__(MCPToolManager)
+        manager.servers_config = {
+            "custom-header-server": {
+                "auth_type": "api_key",
+                "auth_header": "X-Custom-Auth",
+                "url": "http://localhost:8080"
+            }
+        }
+        manager._user_clients = {}
+        manager._user_clients_lock = asyncio.Lock()
+        manager._create_log_handler = MagicMock(return_value=None)
+        manager._create_elicitation_handler = MagicMock(return_value=None)
+        manager._create_sampling_handler = MagicMock(return_value=None)
+
+        with patch("atlas.modules.mcp_tools.token_storage.get_token_storage") as mock_storage, \
+             patch("atlas.modules.mcp_tools.client.Client") as mock_client_class, \
+             patch("atlas.modules.mcp_tools.client.StreamableHttpTransport") as mock_transport_class:
+
+            mock_token = MagicMock()
+            mock_token.token_value = "custom-key-456"
+            mock_token_storage = MagicMock()
+            mock_token_storage.get_valid_token.return_value = mock_token
+            mock_storage.return_value = mock_token_storage
+
+            mock_transport = MagicMock()
+            mock_transport_class.return_value = mock_transport
+            mock_client = MagicMock()
+            mock_client_class.return_value = mock_client
+
+            result = await manager._get_user_client("custom-header-server", "user@example.com")
+
+            assert result is mock_client
+            # Verify custom header name was used
+            transport_call_kwargs = mock_transport_class.call_args
+            assert transport_call_kwargs[1]["headers"] == {"X-Custom-Auth": "custom-key-456"}
+
+    @pytest.mark.asyncio
+    async def test_caches_client(self, manager):
+        """Should cache client for subsequent calls."""
+        with patch("atlas.modules.mcp_tools.token_storage.get_token_storage") as mock_storage, \
+             patch("atlas.modules.mcp_tools.client.Client") as mock_client_class, \
+             patch("atlas.modules.mcp_tools.client.StreamableHttpTransport") as mock_transport_class:
+
+            mock_token = MagicMock()
+            mock_token.token_value = "test-api-key"
+            mock_token_storage = MagicMock()
+            mock_token_storage.get_valid_token.return_value = mock_token
+            mock_storage.return_value = mock_token_storage
+
+            mock_transport = MagicMock()
+            mock_transport_class.return_value = mock_transport
+            mock_client = MagicMock()
+            mock_client_class.return_value = mock_client
+
+            # First call creates client
+            result1 = await manager._get_user_client("test-server", "user@example.com")
+            # Second call should use cache
+            result2 = await manager._get_user_client("test-server", "user@example.com")
+
+            assert result1 is result2
+            # Client constructor only called once
+            assert mock_client_class.call_count == 1
+
+    @pytest.mark.asyncio
+    async def test_invalidates_cache_on_expired_token(self, manager):
+        """Should invalidate cached client when token expires."""
+        with patch("atlas.modules.mcp_tools.token_storage.get_token_storage") as mock_storage, \
+             patch("atlas.modules.mcp_tools.client.Client") as mock_client_class, \
+             patch("atlas.modules.mcp_tools.client.StreamableHttpTransport") as mock_transport_class:
+
+            mock_token = MagicMock()
+            mock_token.token_value = "test-api-key"
+            mock_token_storage = MagicMock()
+            mock_storage.return_value = mock_token_storage
+
+            mock_transport = MagicMock()
+            mock_transport_class.return_value = mock_transport
+            mock_client = MagicMock()
+            mock_client_class.return_value = mock_client
+
+            # First call - token valid
+            mock_token_storage.get_valid_token.return_value = mock_token
+            result1 = await manager._get_user_client("test-server", "user@example.com")
+            assert result1 is mock_client
+
+            # Second call - token expired (returns None)
+            mock_token_storage.get_valid_token.return_value = None
+            result2 = await manager._get_user_client("test-server", "user@example.com")
+
+            # Should return None and cache should be invalidated
+            assert result2 is None
+            cache_key = ("user@example.com", "test-server")
+            assert cache_key not in manager._user_clients
+
+
+class TestInvalidateUserClient:
+    """Test _invalidate_user_client method."""
+
+    @pytest.mark.asyncio
+    async def test_removes_cached_client(self):
+        """Should remove client from cache."""
+        import asyncio
+
+        from atlas.modules.mcp_tools.client import MCPToolManager
+
+        manager = MCPToolManager.__new__(MCPToolManager)
+        manager._user_clients = {
+            ("user@example.com", "test-server"): MagicMock(),
+            ("other@example.com", "test-server"): MagicMock(),
+        }
+        manager._user_clients_lock = asyncio.Lock()
+
+        await manager._invalidate_user_client("user@example.com", "test-server")
+
+        assert ("user@example.com", "test-server") not in manager._user_clients
+        # Other user's client should remain
+        assert ("other@example.com", "test-server") in manager._user_clients
+
+    @pytest.mark.asyncio
+    async def test_handles_missing_cache_entry(self):
+        """Should not error when cache entry doesn't exist."""
+        import asyncio
+
+        from atlas.modules.mcp_tools.client import MCPToolManager
+
+        manager = MCPToolManager.__new__(MCPToolManager)
+        manager._user_clients = {}
+        manager._user_clients_lock = asyncio.Lock()
+
+        # Should not raise
+        await manager._invalidate_user_client("user@example.com", "test-server")

atlas/tests/test_mcp_data_injection.py

@@ -0,0 +1,270 @@
+"""
+Tests for _mcp_data injection into MCP tool arguments.
+
+Validates that tools declaring an _mcp_data parameter in their schema
+receive structured metadata about all available MCP tools, following
+the same pattern as the username injection feature.
+"""
+
+from unittest.mock import MagicMock
+
+from atlas.application.chat.utilities.tool_executor import (
+    build_mcp_data,
+    inject_context_into_args,
+    tool_accepts_mcp_data,
+)
+
+
+class FakeTool:
+    """Minimal tool object matching the MCP tool interface."""
+
+    def __init__(self, name, description="", inputSchema=None):
+        self.name = name
+        self.description = description
+        self.inputSchema = inputSchema or {"type": "object", "properties": {}}
+
+
+def _make_tool_manager(available_tools=None, schema_override=None):
+    """Create a mock tool manager with the given available_tools dict."""
+    manager = MagicMock()
+    manager.available_tools = available_tools or {}
+
+    def get_tools_schema(tool_names):
+        if schema_override is not None:
+            return schema_override
+        schemas = []
+        for server_name, server_data in manager.available_tools.items():
+            for tool in server_data.get("tools", []):
+                fq_name = f"{server_name}_{tool.name}"
+                if fq_name in tool_names:
+                    schemas.append({
+                        "type": "function",
+                        "function": {
+                            "name": fq_name,
+                            "description": tool.description or "",
+                            "parameters": tool.inputSchema or {},
+                        },
+                    })
+        return schemas
+
+    manager.get_tools_schema = MagicMock(side_effect=get_tools_schema)
+    return manager
+
+
+# -- tool_accepts_mcp_data tests --
+
+
+class TestToolAcceptsMcpData:
+    """Tests for tool_accepts_mcp_data detection function."""
+
+    def test_returns_true_when_schema_has_mcp_data(self):
+        tool = FakeTool(
+            "planner",
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "task": {"type": "string"},
+                    "_mcp_data": {"type": "object"},
+                },
+            },
+        )
+        manager = _make_tool_manager({"demo": {"tools": [tool], "config": {}}})
+        assert tool_accepts_mcp_data("demo_planner", manager) is True
+
+    def test_returns_false_when_schema_lacks_mcp_data(self):
+        tool = FakeTool(
+            "search",
+            inputSchema={
+                "type": "object",
+                "properties": {"query": {"type": "string"}},
+            },
+        )
+        manager = _make_tool_manager({"demo": {"tools": [tool], "config": {}}})
+        assert tool_accepts_mcp_data("demo_search", manager) is False
+
+    def test_returns_false_with_no_tool_name(self):
+        assert tool_accepts_mcp_data("", MagicMock()) is False
+
+    def test_returns_false_with_no_tool_manager(self):
+        assert tool_accepts_mcp_data("some_tool", None) is False
+
+    def test_returns_false_when_schema_lookup_fails(self):
+        manager = MagicMock()
+        manager.get_tools_schema = MagicMock(side_effect=RuntimeError("fail"))
+        assert tool_accepts_mcp_data("any_tool", manager) is False
+
+
+# -- build_mcp_data tests --
+
+
+class TestBuildMcpData:
+    """Tests for build_mcp_data output structure."""
+
+    def test_returns_empty_when_no_tools(self):
+        manager = _make_tool_manager({})
+        result = build_mcp_data(manager)
+        assert result == {"available_servers": []}
+
+    def test_returns_empty_when_no_manager(self):
+        result = build_mcp_data(None)
+        assert result == {"available_servers": []}
+
+    def test_skips_canvas_server(self):
+        tool = FakeTool("canvas")
+        manager = _make_tool_manager({
+            "canvas": {"tools": [tool], "config": {}},
+        })
+        result = build_mcp_data(manager)
+        assert result["available_servers"] == []
+
+    def test_includes_server_and_tool_metadata(self):
+        tool_a = FakeTool(
+            "search",
+            description="Search documents",
+            inputSchema={
+                "type": "object",
+                "properties": {"query": {"type": "string"}},
+            },
+        )
+        tool_b = FakeTool("list", description="List items")
+        manager = _make_tool_manager({
+            "myserver": {
+                "tools": [tool_a, tool_b],
+                "config": {"description": "My Server"},
+            },
+        })
+
+        result = build_mcp_data(manager)
+        assert len(result["available_servers"]) == 1
+        server = result["available_servers"][0]
+        assert server["server_name"] == "myserver"
+        assert server["description"] == "My Server"
+        assert len(server["tools"]) == 2
+
+        tool_entry = server["tools"][0]
+        assert tool_entry["name"] == "myserver_search"
+        assert tool_entry["description"] == "Search documents"
+        assert "properties" in tool_entry["parameters"]
+
+    def test_multiple_servers(self):
+        tool1 = FakeTool("t1")
+        tool2 = FakeTool("t2")
+        manager = _make_tool_manager({
+            "server_a": {"tools": [tool1], "config": {}},
+            "server_b": {"tools": [tool2], "config": {}},
+        })
+        result = build_mcp_data(manager)
+        names = [s["server_name"] for s in result["available_servers"]]
+        assert "server_a" in names
+        assert "server_b" in names
+
+    def test_handles_missing_description(self):
+        tool = FakeTool("t", description=None, inputSchema=None)
+        manager = _make_tool_manager({
+            "s": {"tools": [tool], "config": {}},
+        })
+        result = build_mcp_data(manager)
+        server = result["available_servers"][0]
+        assert server["description"] == ""
+        assert server["tools"][0]["description"] == ""
+        assert server["tools"][0]["parameters"] == {"type": "object", "properties": {}}
+
+    def test_uses_short_description_fallback(self):
+        tool = FakeTool("t")
+        manager = _make_tool_manager({
+            "s": {"tools": [tool], "config": {"short_description": "Short desc"}},
+        })
+        result = build_mcp_data(manager)
+        assert result["available_servers"][0]["description"] == "Short desc"
+
+
+# -- inject_context_into_args with _mcp_data tests --
+
+
+class TestInjectMcpData:
+    """Tests for _mcp_data injection in inject_context_into_args."""
+
+    def test_injects_mcp_data_when_tool_accepts_it(self):
+        tool = FakeTool(
+            "planner",
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "task": {"type": "string"},
+                    "_mcp_data": {"type": "object"},
+                },
+            },
+        )
+        manager = _make_tool_manager({
+            "demo": {"tools": [tool], "config": {"description": "Demo"}},
+        })
+
+        result = inject_context_into_args(
+            {"task": "do something"},
+            {"user_email": "user@test.com"},
+            "demo_planner",
+            manager,
+        )
+
+        assert "_mcp_data" in result
+        assert "available_servers" in result["_mcp_data"]
+        assert len(result["_mcp_data"]["available_servers"]) == 1
+
+    def test_does_not_inject_mcp_data_when_tool_lacks_param(self):
+        tool = FakeTool(
+            "search",
+            inputSchema={
+                "type": "object",
+                "properties": {"query": {"type": "string"}},
+            },
+        )
+        manager = _make_tool_manager({
+            "demo": {"tools": [tool], "config": {}},
+        })
+
+        result = inject_context_into_args(
+            {"query": "hello"},
+            {"user_email": "user@test.com"},
+            "demo_search",
+            manager,
+        )
+
+        assert "_mcp_data" not in result
+
+    def test_mcp_data_reinjected_after_edit(self):
+        """Simulates the re-injection path after user edits tool arguments."""
+        tool = FakeTool(
+            "planner",
+            inputSchema={
+                "type": "object",
+                "properties": {
+                    "task": {"type": "string"},
+                    "_mcp_data": {"type": "object"},
+                },
+            },
+        )
+        manager = _make_tool_manager({
+            "demo": {"tools": [tool], "config": {}},
+        })
+
+        # Simulate user editing args (removing _mcp_data)
+        edited_args = {"task": "edited task"}
+        result = inject_context_into_args(
+            edited_args,
+            {"user_email": "user@test.com"},
+            "demo_planner",
+            manager,
+        )
+
+        # _mcp_data should be re-injected
+        assert "_mcp_data" in result
+        assert result["task"] == "edited task"
+
+    def test_mcp_data_not_injected_without_tool_manager(self):
+        result = inject_context_into_args(
+            {"task": "test"},
+            {"user_email": "user@test.com"},
+            "some_tool",
+            None,
+        )
+        assert "_mcp_data" not in result

atlas/tests/test_mcp_get_authorized_servers.py

@@ -0,0 +1,95 @@
+"""Test get_authorized_servers with async auth function."""
+
+import pytest
+
+from atlas.modules.mcp_tools.client import MCPToolManager
+
+
+@pytest.mark.asyncio
+async def test_get_authorized_servers_with_async_auth():
+    """Test that get_authorized_servers properly handles async auth_check_func."""
+
+    # Create a mock MCPToolManager with test server config
+    mcp_manager = MCPToolManager(None)
+    mcp_manager.servers_config = {
+        "server1": {
+            "enabled": True,
+            "groups": ["admin", "users"]
+        },
+        "server2": {
+            "enabled": True,
+            "groups": ["admin"]
+        },
+        "server3": {
+            "enabled": True,
+            "groups": []  # No groups required
+        },
+        "server4": {
+            "enabled": False,
+            "groups": ["admin"]
+        }
+    }
+
+    # Mock async auth function
+    async def mock_auth_check(user_email: str, group: str) -> bool:
+        """Mock auth check that returns True for admin group."""
+        return group == "admin"
+
+    # Test with user who has admin access
+    authorized = await mcp_manager.get_authorized_servers("admin@test.com", mock_auth_check)
+
+    # Should include server1 (has admin), server2 (has admin), server3 (no groups required)
+    # Should NOT include server4 (disabled)
+    assert set(authorized) == {"server1", "server2", "server3"}
+
+
+@pytest.mark.asyncio
+async def test_get_authorized_servers_with_multiple_groups():
+    """Test authorization with multiple group checks."""
+
+    mcp_manager = MCPToolManager(None)
+    mcp_manager.servers_config = {
+        "server1": {
+            "enabled": True,
+            "groups": ["users", "developers"]
+        },
+        "server2": {
+            "enabled": True,
+            "groups": ["admin"]
+        }
+    }
+
+    # User is in 'users' group but not 'admin'
+    async def mock_auth_check(user_email: str, group: str) -> bool:
+        return group in ["users", "developers"]
+
+    authorized = await mcp_manager.get_authorized_servers("user@test.com", mock_auth_check)
+
+    # Should include server1 (user is in 'users' group)
+    # Should NOT include server2 (user not in 'admin' group)
+    assert authorized == ["server1"]
+
+
+@pytest.mark.asyncio
+async def test_get_authorized_servers_no_access():
+    """Test when user has no access to any servers."""
+
+    mcp_manager = MCPToolManager(None)
+    mcp_manager.servers_config = {
+        "server1": {
+            "enabled": True,
+            "groups": ["admin"]
+        },
+        "server2": {
+            "enabled": True,
+            "groups": ["superusers"]
+        }
+    }
+
+    # User has no group memberships
+    async def mock_auth_check(user_email: str, group: str) -> bool:
+        return False
+
+    authorized = await mcp_manager.get_authorized_servers("user@test.com", mock_auth_check)
+
+    assert authorized == []