atlas-chat 0.1.0__py3-none-any.whl

atlas/routes/mcp_auth_routes.py

@@ -0,0 +1,223 @@
+"""MCP Authentication routes for per-user API key and token management.
+
+This module provides user-facing endpoints for managing authentication tokens
+with MCP servers. Each user's tokens are stored separately and securely encrypted.
+
+Users can manually upload API keys, JWTs, or bearer tokens for MCP servers that
+require authentication. This supports any type of token that can be passed as a
+bearer token in the Authorization header.
+
+Updated: 2025-01-21
+"""
+
+import logging
+from typing import Optional
+
+from fastapi import APIRouter, Depends, HTTPException
+from pydantic import BaseModel
+
+from atlas.core.auth import is_user_in_group
+from atlas.core.log_sanitizer import get_current_user, sanitize_for_logging
+from atlas.infrastructure.app_factory import app_factory
+from atlas.modules.mcp_tools.token_storage import get_token_storage
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter(prefix="/api/mcp/auth", tags=["mcp-auth"])
+
+
+class TokenUpload(BaseModel):
+    """Request body for uploading an API key or token."""
+    token: str
+    expires_at: Optional[float] = None  # Unix timestamp, or None for no expiry
+    scopes: Optional[str] = None  # Space-separated scopes (optional)
+
+
+# --- Routes ---
+
+
+@router.get("/status")
+async def get_auth_status(current_user: str = Depends(get_current_user)):
+    """Get authentication status for all MCP servers accessible to the user.
+
+    Returns information about which servers require authentication,
+    which the user has authenticated with, and token status.
+    """
+    try:
+        mcp_manager = app_factory.get_mcp_manager()
+        token_storage = get_token_storage()
+
+        # Get servers the user is authorized to access
+        authorized_servers = await mcp_manager.get_authorized_servers(
+            current_user, is_user_in_group
+        )
+
+        # Get user's current auth status
+        user_auth_status = token_storage.get_user_auth_status(current_user)
+
+        # Build response with server auth requirements and user's status
+        servers_status = []
+
+        for server_name in authorized_servers:
+            server_config = mcp_manager.servers_config.get(server_name, {})
+            auth_type = server_config.get("auth_type", "none")
+
+            # Get user's token status for this server
+            token_status = user_auth_status.get(server_name)
+
+            server_info = {
+                "server_name": server_name,
+                "auth_type": auth_type,
+                "auth_required": auth_type != "none",
+                "authenticated": token_status is not None,
+                "description": server_config.get("description", ""),
+            }
+
+            # Add token details if authenticated
+            if token_status:
+                server_info.update({
+                    "token_type": token_status["token_type"],
+                    "is_expired": token_status["is_expired"],
+                    "expires_at": token_status["expires_at"],
+                    "time_until_expiry": token_status["time_until_expiry"],
+                    "has_refresh_token": token_status["has_refresh_token"],
+                    "scopes": token_status["scopes"],
+                })
+
+            servers_status.append(server_info)
+
+        return {
+            "servers": servers_status,
+            "user": current_user,
+        }
+
+    except Exception as e:
+        logger.error(f"Error getting auth status: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail="Internal server error while fetching auth status")
+
+
+@router.post("/{server_name}/token")
+async def upload_token(
+    server_name: str,
+    token_data: TokenUpload,
+    current_user: str = Depends(get_current_user),
+):
+    """Upload an API key or bearer token for an MCP server.
+
+    This allows users to manually provide tokens for servers that require
+    authentication. Tokens can be:
+    - API keys
+    - JWT tokens
+    - Bearer tokens
+    - Any other string token that can be used in Authorization header
+
+    The token will be securely encrypted and stored per-user.
+    """
+    try:
+        mcp_manager = app_factory.get_mcp_manager()
+        token_storage = get_token_storage()
+
+        # Verify server exists and user has access
+        authorized_servers = await mcp_manager.get_authorized_servers(
+            current_user, is_user_in_group
+        )
+
+        if server_name not in authorized_servers:
+            raise HTTPException(
+                status_code=403,
+                detail=f"Not authorized to access server '{server_name}'"
+            )
+
+        # Verify server accepts token authentication
+        server_config = mcp_manager.servers_config.get(server_name, {})
+        auth_type = server_config.get("auth_type", "none")
+
+        if auth_type not in ("jwt", "bearer", "api_key", "oauth"):
+            raise HTTPException(
+                status_code=400,
+                detail=f"Server '{server_name}' does not accept token authentication (auth_type: {auth_type})"
+            )
+
+        # Validate token is not empty
+        if not token_data.token or not token_data.token.strip():
+            raise HTTPException(
+                status_code=400,
+                detail="Token cannot be empty"
+            )
+
+        # Determine token type based on server config
+        token_type = "bearer"
+        if auth_type == "jwt":
+            token_type = "jwt"
+        elif auth_type == "api_key":
+            token_type = "api_key"
+
+        # Store the token
+        stored_token = token_storage.store_token(
+            user_email=current_user,
+            server_name=server_name,
+            token_value=token_data.token.strip(),
+            token_type=token_type,
+            expires_at=token_data.expires_at,
+            scopes=token_data.scopes,
+        )
+
+        sanitized_server = sanitize_for_logging(server_name)
+        logger.info(f"User uploaded token for MCP server '{sanitized_server}'")
+
+        return {
+            "message": f"Token stored for server '{server_name}'",
+            "server_name": server_name,
+            "token_type": stored_token.token_type,
+            "expires_at": stored_token.expires_at,
+            "scopes": stored_token.scopes,
+        }
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Error uploading token: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail="Internal server error while uploading token")
+
+
+@router.delete("/{server_name}/token")
+async def remove_token(
+    server_name: str,
+    current_user: str = Depends(get_current_user),
+):
+    """Remove stored token for an MCP server (disconnect).
+
+    This removes the user's authentication token for the specified server.
+    The user will need to re-authenticate to use the server's tools.
+    """
+    try:
+        token_storage = get_token_storage()
+
+        # Remove the token
+        removed = token_storage.remove_token(current_user, server_name)
+
+        if not removed:
+            raise HTTPException(
+                status_code=404,
+                detail=f"No token found for server '{server_name}'"
+            )
+
+        # Invalidate any cached client for this user/server combination
+        tool_manager = app_factory.get_mcp_manager()
+        if tool_manager is not None:
+            await tool_manager._invalidate_user_client(current_user, server_name)
+            logger.debug(f"Invalidated cached client for server '{server_name}'")
+
+        sanitized_server = sanitize_for_logging(server_name)
+        logger.info(f"User removed token for MCP server '{sanitized_server}'")
+
+        return {
+            "message": f"Token removed for server '{server_name}'",
+            "server_name": server_name,
+        }
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Error removing token: {e}", exc_info=True)
+        raise HTTPException(status_code=500, detail="Internal server error while removing token")

atlas/server_cli.py

@@ -0,0 +1,164 @@
+"""
+Atlas Server CLI - Start the Atlas backend server.
+
+Usage:
+    atlas-server                              # Start with defaults
+    atlas-server --port 8000                  # Custom port
+    atlas-server --env /path/to/.env          # Custom env file
+    atlas-server --config-folder /path/to/config  # Custom config folder
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+
+def _apply_env_file(env_path: Path) -> None:
+    """Load environment variables from a .env file."""
+    from dotenv import load_dotenv
+
+    if not env_path.exists():
+        print(f"Error: env file not found: {env_path}", file=sys.stderr)
+        sys.exit(2)
+
+    load_dotenv(dotenv_path=str(env_path))
+
+
+def _apply_config_folder(config_folder: Path) -> None:
+    """Set up config folder as the overrides directory."""
+    if not config_folder.exists():
+        print(f"Error: config folder not found: {config_folder}", file=sys.stderr)
+        sys.exit(2)
+
+    os.environ["APP_CONFIG_OVERRIDES"] = str(config_folder.resolve())
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """Build argument parser for atlas-server CLI."""
+    parser = argparse.ArgumentParser(
+        prog="atlas-server",
+        description="Start the Atlas backend server with MCP integration.",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=None,
+        help="Port to run the server on (default: 8000 or PORT env var).",
+    )
+    parser.add_argument(
+        "--host",
+        default=None,
+        help="Host to bind to (default: 127.0.0.1 or ATLAS_HOST env var).",
+    )
+    parser.add_argument(
+        "--env",
+        dest="env_file",
+        default=None,
+        help="Path to .env file (default: .env in current directory or package root).",
+    )
+    parser.add_argument(
+        "--config-folder",
+        dest="config_folder",
+        default=None,
+        help="Path to config folder for overrides (sets APP_CONFIG_OVERRIDES).",
+    )
+    parser.add_argument(
+        "--config-defaults",
+        dest="config_defaults",
+        default=None,
+        help="Path to config defaults folder (sets APP_CONFIG_DEFAULTS).",
+    )
+    parser.add_argument(
+        "--reload",
+        action="store_true",
+        help="Enable auto-reload for development (not recommended for production).",
+    )
+    parser.add_argument(
+        "--workers",
+        type=int,
+        default=1,
+        help="Number of worker processes (default: 1).",
+    )
+    parser.add_argument(
+        "--version",
+        action="store_true",
+        help="Print version and exit.",
+    )
+    return parser
+
+
+def run_server(args: argparse.Namespace) -> int:
+    """Run the Atlas server with the given arguments."""
+    import uvicorn
+
+    # Determine host and port
+    host = args.host or os.getenv("ATLAS_HOST", "127.0.0.1")
+    port = args.port or int(os.getenv("PORT", "8000"))
+
+    # Import the FastAPI app
+    from atlas.main import app
+
+    print(f"Starting Atlas server on {host}:{port}")
+
+    if args.reload:
+        print("Warning: --reload is enabled. This is not recommended for production.")
+        uvicorn.run(
+            "atlas.main:app",
+            host=host,
+            port=port,
+            reload=True,
+            workers=1,  # reload doesn't support multiple workers
+        )
+    else:
+        uvicorn.run(
+            app,
+            host=host,
+            port=port,
+            workers=args.workers,
+        )
+
+    return 0
+
+
+def main() -> None:
+    """Main entry point for atlas-server CLI."""
+    parser = build_parser()
+    args = parser.parse_args()
+
+    if args.version:
+        from atlas.version import VERSION
+        print(f"atlas-server version {VERSION}")
+        sys.exit(0)
+
+    # Apply env file first (before any other imports that might use env vars)
+    if args.env_file:
+        _apply_env_file(Path(args.env_file).expanduser())
+    else:
+        # Try to find .env in current directory or package root
+        cwd_env = Path.cwd() / ".env"
+        if cwd_env.exists():
+            _apply_env_file(cwd_env)
+        else:
+            # Check package root (parent of atlas/)
+            pkg_root_env = Path(__file__).resolve().parents[1] / ".env"
+            if pkg_root_env.exists():
+                _apply_env_file(pkg_root_env)
+
+    # Apply config folder if specified
+    if args.config_folder:
+        _apply_config_folder(Path(args.config_folder).expanduser())
+
+    # Apply config defaults if specified
+    if args.config_defaults:
+        defaults_path = Path(args.config_defaults).expanduser()
+        if not defaults_path.exists():
+            print(f"Error: config defaults folder not found: {defaults_path}", file=sys.stderr)
+            sys.exit(2)
+        os.environ["APP_CONFIG_DEFAULTS"] = str(defaults_path.resolve())
+
+    sys.exit(run_server(args))
+
+
+if __name__ == "__main__":
+    main()

atlas/tests/conftest.py

@@ -0,0 +1,20 @@
+import sys
+from pathlib import Path
+
+# Ensure the backend root is on sys.path for absolute imports like 'infrastructure.*'
+backend_root = Path(__file__).resolve().parents[1]
+project_root = backend_root.parent
+if str(backend_root) not in sys.path:
+    sys.path.insert(0, str(backend_root))
+if str(project_root) not in sys.path:
+    sys.path.insert(0, str(project_root))
+
+# Pre-import critical modules before any test files can replace them with fakes.
+# This prevents test pollution where one test file patches sys.modules and other
+# tests import the fake instead of the real module.
+# See test_capability_tokens_and_injection.py which patches LiteLLMCaller.
+import atlas.modules.llm.litellm_caller  # noqa: E402, F401
+
+# Explicitly reference the module to satisfy static analyzers that flag unused imports.
+# The import above is intentional: it pre-populates sys.modules with the real module.
+_ = atlas.modules.llm.litellm_caller.LiteLLMCaller  # noqa: E402

atlas/tests/integration/test_mcp_auth_integration.py

@@ -0,0 +1,152 @@
+from unittest.mock import ANY, AsyncMock, Mock, patch
+
+import pytest
+
+from atlas.modules.mcp_tools.client import MCPToolManager
+
+
+@pytest.mark.integration
+class TestMCPAuthenticationIntegration:
+    """Integration tests for MCP authentication."""
+
+    @pytest.mark.asyncio
+    async def test_authenticated_connection_success(self, monkeypatch):
+        """Should successfully connect to authenticated MCP server."""
+        monkeypatch.setenv("MCP_TEST_TOKEN", "test-api-key-123")
+
+        # Configure test server with auth requirement
+        server_config = {
+            "url": "http://localhost:8001/mcp",
+            "transport": "http",
+            "auth_token": "${MCP_TEST_TOKEN}"
+        }
+
+        # Mock config_manager to return our test server config
+        with patch('atlas.modules.mcp_tools.client.config_manager') as mock_config_manager:
+            mock_config_manager.mcp_config.servers = {"mcp-http-mock": Mock()}
+            mock_config_manager.mcp_config.servers["mcp-http-mock"].model_dump.return_value = server_config
+            mock_config_manager.app_settings.mcp_call_timeout = 120
+            mock_config_manager.app_settings.mcp_discovery_timeout = 30
+
+            manager = MCPToolManager()
+            manager.servers_config = {"mcp-http-mock": server_config}
+
+            # Mock the fastmcp.Client to avoid actual network call for now
+            with patch('atlas.modules.mcp_tools.client.Client') as MockFastMCPClient:
+                mock_client_instance = MockFastMCPClient.return_value
+                mock_client_instance.__aenter__.return_value = mock_client_instance
+                mock_client_instance.list_tools.return_value = [] # Mock an empty list of tools
+
+                await manager.initialize_clients()
+
+                # Assert that the client was initialized and added to manager.clients
+                assert "mcp-http-mock" in manager.clients
+                MockFastMCPClient.assert_called_once_with(
+                    "http://localhost:8001/mcp",
+                    auth="test-api-key-123",
+                    log_handler=ANY,
+                    elicitation_handler=ANY,
+                    sampling_handler=ANY,
+                )
+
+    @pytest.mark.asyncio
+    async def test_authenticated_connection_failure_invalid_token(self, monkeypatch):
+        """Should fail to connect with invalid token."""
+        monkeypatch.setenv("MCP_TEST_TOKEN", "invalid-token") # Set an invalid token
+
+        server_config = {
+            "url": "http://localhost:8001/mcp",
+            "transport": "http",
+            "auth_token": "${MCP_TEST_TOKEN}"
+        }
+
+        with patch('atlas.modules.mcp_tools.client.config_manager') as mock_config_manager:
+            mock_config_manager.mcp_config.servers = {"mcp-http-mock": Mock()}
+            mock_config_manager.mcp_config.servers["mcp-http-mock"].model_dump.return_value = server_config
+            mock_config_manager.app_settings.mcp_call_timeout = 120
+            mock_config_manager.app_settings.mcp_discovery_timeout = 30
+
+            manager = MCPToolManager()
+            manager.servers_config = {"mcp-http-mock": server_config}
+
+            # Mock the fastmcp.Client to simulate an authentication error
+            with patch('atlas.modules.mcp_tools.client.Client') as MockFastMCPClient:
+                mock_client_instance = MockFastMCPClient.return_value
+                mock_client_instance.__aenter__.side_effect = Exception("Authentication failed: 401 Unauthorized")
+
+                await manager.initialize_clients()
+
+                # Client object is created successfully (not connected yet)
+                assert "mcp-http-mock" in manager.clients
+                # Verify auth token was passed correctly
+                MockFastMCPClient.assert_called_once_with(
+                    "http://localhost:8001/mcp",
+                    auth="invalid-token",
+                    log_handler=ANY,
+                    elicitation_handler=ANY,
+                    sampling_handler=ANY,
+                )
+
+                # Now try to discover tools - this should fail due to auth error
+                await manager.discover_tools()
+
+                # After failed connection, tools should not be discovered
+                if not hasattr(manager, '_tool_index'):
+                    manager._tool_index = {}
+                assert len([k for k in manager._tool_index.keys() if k.startswith("mcp-http-mock_")]) == 0
+
+    @pytest.mark.asyncio
+    async def test_authenticated_tool_execution(self, monkeypatch):
+        """Should successfully execute tool after authentication."""
+        monkeypatch.setenv("MCP_TEST_TOKEN", "test-api-key-123")
+
+        server_config = {
+            "url": "http://localhost:8001/mcp",
+            "transport": "http",
+            "auth_token": "${MCP_TEST_TOKEN}"
+        }
+
+        with patch('atlas.modules.mcp_tools.client.config_manager') as mock_config_manager:
+            mock_config_manager.mcp_config.servers = {"mcp-http-mock": Mock()}
+            mock_config_manager.mcp_config.servers["mcp-http-mock"].model_dump.return_value = server_config
+            mock_config_manager.app_settings.mcp_call_timeout = 120
+            mock_config_manager.app_settings.mcp_discovery_timeout = 30
+
+            manager = MCPToolManager()
+            manager.servers_config = {"mcp-http-mock": server_config}
+
+            # Mock the fastmcp.Client and its call_tool method
+            with patch('atlas.modules.mcp_tools.client.Client') as MockFastMCPClient:
+                mock_client_instance = MockFastMCPClient.return_value
+                mock_client_instance.__aenter__.return_value = mock_client_instance
+                mock_client_instance.list_tools.return_value = [] # Mock an empty list of tools
+
+                # Make call_tool return an async result
+                class MockResult:
+                    data = {"results": "tool_result"}
+                mock_client_instance.call_tool = AsyncMock(return_value=MockResult())
+
+                await manager.initialize_clients()
+
+                # Simulate tool call
+                from atlas.domain.messages.models import ToolCall
+                tool_call = ToolCall(id="call_1", name="mcp-http-mock_test_tool", arguments={})
+
+                # Need to mock the _tool_index for execute_tool to find the server
+                mock_tool = Mock()
+                mock_tool.name = "test_tool"
+                manager._tool_index = {
+                    "mcp-http-mock_test_tool": {
+                        'server': 'mcp-http-mock',
+                        'tool': mock_tool
+                    }
+                }
+
+                result = await manager.execute_tool(tool_call)
+
+                assert result.success is True
+                assert "tool_result" in result.content
+                mock_client_instance.call_tool.assert_called_once()
+                call_args = mock_client_instance.call_tool.call_args
+                assert call_args[0] == ("test_tool", {})
+                assert "progress_handler" in call_args[1]

atlas/tests/manual_test_sampling.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+"""
+Quick test script to verify sampling functionality works end-to-end.
+This tests the sampling_demo server directly to ensure it can be initialized
+and sampling works correctly.
+"""
+
+import asyncio
+import logging
+import sys
+from pathlib import Path
+
+# Add backend to path
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+async def test_sampling_demo():
+    """Test the sampling_demo server."""
+    try:
+        from fastmcp import Client
+
+        logger.info("Starting sampling demo test...")
+
+        # Create sampling handler that uses LiteLLM
+        async def sampling_handler(messages, params, context):
+            from mcp.types import CreateMessageResult, TextContent
+
+            logger.info(f"Sampling handler called with {len(messages)} messages")
+
+            # For testing purposes, we'll return a mock response
+            # In production, this would call the actual LLM
+            logger.info("Returning mock LLM response for testing")
+
+            # Return proper CreateMessageResult
+            return CreateMessageResult(
+                role="assistant",
+                content=TextContent(type="text", text="This is a mock LLM response for testing purposes."),
+                model="mock-model"
+            )
+
+        # Create client to sampling_demo server
+        server_path = "backend/mcp/sampling_demo/main.py"
+        logger.info(f"Connecting to server: {server_path}")
+
+        client = Client(server_path, sampling_handler=sampling_handler)
+
+        async with client:
+            # List tools
+            tools = await client.list_tools()
+            logger.info(f"Available tools: {[t.name for t in tools]}")
+
+            # Test summarize_text tool
+            logger.info("\n=== Testing summarize_text ===")
+            result = await client.call_tool(
+                "summarize_text",
+                {
+                    "text": "The quick brown fox jumps over the lazy dog. "
+                           "This sentence contains every letter of the alphabet. "
+                           "It is commonly used for testing fonts and keyboards."
+                }
+            )
+            logger.info(f"Result: {result}")
+
+            # Test analyze_sentiment tool
+            logger.info("\n=== Testing analyze_sentiment ===")
+            result = await client.call_tool(
+                "analyze_sentiment",
+                {
+                    "text": "I absolutely love this product! It's amazing and works perfectly!"
+                }
+            )
+            logger.info(f"Result: {result}")
+
+        logger.info("\n✅ All tests passed!")
+        return True
+
+    except Exception as e:
+        logger.error(f"Test failed: {e}", exc_info=True)
+        return False
+
+
+if __name__ == "__main__":
+    success = asyncio.run(test_sampling_demo())
+    sys.exit(0 if success else 1)