atlas-chat 0.1.0__py3-none-any.whl

atlas/tests/test_log_level_sensitive_data.py

@@ -0,0 +1,285 @@
+"""Tests for log level control of sensitive data logging.
+
+These tests verify that:
+1. User message content is only logged at DEBUG level, not INFO
+2. LLM response content is only logged at DEBUG level, not INFO
+3. Non-sensitive metadata is always logged at INFO level
+4. The LOG_LEVEL environment variable controls this behavior
+"""
+
+import asyncio
+import importlib
+import logging
+import sys
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from atlas.application.chat.service import ChatService
+
+
+def _ensure_real_litellm_module():
+    """
+    Ensure the real litellm_caller module is in sys.modules.
+
+    Some test files (e.g., test_capability_tokens_and_injection.py) patch the
+    litellm_caller module at import time. This function forces a reimport of
+    the real module to ensure patches work correctly.
+    """
+    module_name = "atlas.modules.llm.litellm_caller"
+
+    # Check if current module is fake
+    if module_name in sys.modules:
+        current_module = sys.modules[module_name]
+        if hasattr(current_module, "LiteLLMCaller"):
+            caller_class = current_module.LiteLLMCaller
+            if not hasattr(caller_class, "_get_model_kwargs"):
+                # It's a fake, remove it and reimport
+                del sys.modules[module_name]
+                importlib.import_module(module_name)
+    else:
+        # Not in sys.modules, import fresh
+        importlib.import_module(module_name)
+
+
+# Module-level LiteLLMCaller variable, set by fixture
+LiteLLMCaller = None
+
+
+@pytest.fixture(autouse=True)
+def ensure_real_litellm_for_tests():
+    """Fixture to ensure real LiteLLM module is loaded before each test."""
+    _ensure_real_litellm_module()
+    # Re-import the class from the now-correct module
+    from atlas.modules.llm.litellm_caller import LiteLLMCaller as RealCaller
+    # Make it available globally for this module
+    global LiteLLMCaller
+    LiteLLMCaller = RealCaller
+    yield
+
+
+# Initially try to get the real class
+_ensure_real_litellm_module()
+_initial_module = importlib.import_module("atlas.modules.llm.litellm_caller")
+LiteLLMCaller = _initial_module.LiteLLMCaller
+
+
+class TestLogLevelSensitiveData:
+    """Tests for log level control of sensitive data."""
+
+    def test_tool_approval_response_summary_excludes_argument_values(self):
+        """Tool approval payloads must never log raw tool arguments at INFO."""
+        from atlas.core.log_sanitizer import summarize_tool_approval_response_for_logging
+
+        payload = {
+            "type": "tool_approval_response",
+            "tool_call_id": "fc_07785773-7583-4e97-bb7d-a24f1f8a0c4b",
+            "approved": True,
+            "arguments": {
+                "file_name": "New_Mexico_Snakes",
+                "markdown_content": "# Snakes of New Mexico\n- secret stuff"
+            },
+            "reason": "User approved"
+        }
+
+        summary = summarize_tool_approval_response_for_logging(payload)
+
+        # Must include safe metadata
+        assert "type=tool_approval_response" in summary
+        assert "tool_call_id=fc_07785773-7583-4e97-bb7d-a24f1f8a0c4b" in summary
+        assert "approved=True" in summary
+
+        # Must not include argument values or reason contents
+        assert "New_Mexico_Snakes" not in summary
+        assert "Snakes of New Mexico" not in summary
+        assert "User approved" not in summary
+
+    def test_chat_service_info_level_excludes_content(self, caplog):
+        """Test that INFO level logging excludes user message content."""
+        # Create mock dependencies
+        mock_llm = MagicMock()
+        mock_tool_manager = MagicMock()
+        mock_connection = MagicMock()
+        mock_config = MagicMock()
+        mock_session_repo = MagicMock()
+        mock_session_repo.get = AsyncMock(return_value=None)
+        mock_session_repo.create = AsyncMock(return_value=None)
+
+        service = ChatService(
+            llm=mock_llm,
+            tool_manager=mock_tool_manager,
+            connection=mock_connection,
+            config_manager=mock_config,
+            session_repository=mock_session_repo
+        )
+
+        # Set log level to INFO
+        with caplog.at_level(logging.INFO):
+            # Create test session first
+            asyncio.run(service.create_session("test-session", "test@test.com"))
+
+            # Clear logs from session creation
+            caplog.clear()
+
+            # Try to call handle_chat_message (it will fail but we only care about logs)
+            try:
+                asyncio.run(service.handle_chat_message(
+                    session_id="test-session",
+                    content="This is sensitive user input that should not be logged at INFO level",
+                    model="test-model",
+                    user_email="test@test.com"
+                ))
+            except Exception:
+                pass  # We expect this to fail, we're just checking logs
+
+        # Check that logs exist but don't contain the sensitive content
+        log_messages = [record.message for record in caplog.records if record.levelno == logging.INFO]
+
+        # Should have INFO log about the call
+        assert any("handle_chat_message called" in msg for msg in log_messages), \
+            "Should have INFO log about handle_chat_message call"
+
+        # Should NOT contain the sensitive content at INFO level
+        assert not any("sensitive user input" in msg for msg in log_messages), \
+            "Should NOT log sensitive content at INFO level"
+
+        # Should log metadata like content length
+        assert any("content_length" in msg for msg in log_messages), \
+            "Should log content_length metadata at INFO level"
+
+    def test_chat_service_debug_level_includes_content(self, caplog):
+        """Test that DEBUG level logging includes user message content."""
+        # Create mock dependencies
+        mock_llm = MagicMock()
+        mock_tool_manager = MagicMock()
+        mock_connection = MagicMock()
+        mock_config = MagicMock()
+        mock_session_repo = MagicMock()
+        mock_session_repo.get = AsyncMock(return_value=None)
+        mock_session_repo.create = AsyncMock(return_value=None)
+
+        service = ChatService(
+            llm=mock_llm,
+            tool_manager=mock_tool_manager,
+            connection=mock_connection,
+            config_manager=mock_config,
+            session_repository=mock_session_repo
+        )
+
+        # Set log level to DEBUG
+        with caplog.at_level(logging.DEBUG):
+            # Create test session first
+            asyncio.run(service.create_session("test-session", "test@test.com"))
+
+            # Clear logs from session creation
+            caplog.clear()
+
+            # Try to call handle_chat_message
+            try:
+                asyncio.run(service.handle_chat_message(
+                    session_id="test-session",
+                    content="This is sensitive user input",
+                    model="test-model",
+                    user_email="test@test.com"
+                ))
+            except Exception:
+                pass  # We expect this to fail, we're just checking logs
+
+        # Check that DEBUG logs include the content
+        log_messages = [record.message for record in caplog.records if record.levelno == logging.DEBUG]
+
+        # Should contain the sensitive content at DEBUG level
+        assert any("sensitive user input" in msg for msg in log_messages), \
+            "Should log sensitive content at DEBUG level"
+
+    @pytest.mark.asyncio
+    async def test_llm_caller_info_level_excludes_response_preview(self, caplog):
+        """Test that INFO level logging excludes LLM response previews."""
+        # Mock the acompletion call
+        mock_response = MagicMock()
+        mock_response.choices = [MagicMock()]
+        mock_response.choices[0].message = MagicMock()
+        mock_response.choices[0].message.content = "This is a sensitive LLM response with user data"
+
+        with patch('atlas.modules.llm.litellm_caller.acompletion', return_value=mock_response):
+            caller = LiteLLMCaller()
+
+            # Mock the config to return test model
+            with patch.object(caller, '_get_litellm_model_name', return_value='gpt-4'):
+                with patch.object(caller, '_get_model_kwargs', return_value={}):
+                    # Set log level to INFO
+                    with caplog.at_level(logging.INFO):
+                        await caller.call_plain(
+                            model_name="test-model",
+                            messages=[{"role": "user", "content": "test"}],
+                            temperature=0.7
+                        )
+
+        # Check logs
+        log_messages = [record.message for record in caplog.records if record.levelno == logging.INFO]
+
+        # Should have INFO log about the call
+        assert any("Plain LLM call" in msg for msg in log_messages), \
+            "Should have INFO log about LLM call"
+
+        # Should NOT contain response preview at INFO level
+        assert not any("sensitive LLM response" in msg for msg in log_messages), \
+            "Should NOT log response preview at INFO level"
+
+        # Should log response length instead
+        assert any("response length" in msg for msg in log_messages), \
+            "Should log response length at INFO level"
+
+    @pytest.mark.asyncio
+    async def test_llm_caller_debug_level_includes_response_preview(self, caplog):
+        """Test that DEBUG level logging includes LLM response previews."""
+        # Mock the acompletion call
+        mock_response = MagicMock()
+        mock_response.choices = [MagicMock()]
+        mock_response.choices[0].message = MagicMock()
+        mock_response.choices[0].message.content = "This is a sensitive LLM response"
+
+        with patch('atlas.modules.llm.litellm_caller.acompletion', return_value=mock_response):
+            caller = LiteLLMCaller()
+
+            # Mock the config to return test model
+            with patch.object(caller, '_get_litellm_model_name', return_value='gpt-4'):
+                with patch.object(caller, '_get_model_kwargs', return_value={}):
+                    # Set log level to DEBUG
+                    with caplog.at_level(logging.DEBUG):
+                        await caller.call_plain(
+                            model_name="test-model",
+                            messages=[{"role": "user", "content": "test"}],
+                            temperature=0.7
+                        )
+
+        # Check logs
+        log_messages = [record.message for record in caplog.records if record.levelno == logging.DEBUG]
+
+        # Should contain response preview at DEBUG level
+        assert any("sensitive LLM response" in msg for msg in log_messages), \
+            "Should log response preview at DEBUG level"
+
+    def test_log_level_from_config_manager(self):
+        """Test that LOG_LEVEL configuration mechanism exists and is functional."""
+        # This test verifies the log level configuration mechanism exists
+        # The actual value will be whatever was set during module initialization
+        from atlas.core.otel_config import OpenTelemetryConfig
+        from atlas.modules.config.config_manager import AppSettings
+
+        # Verify AppSettings has log_level field
+        app_settings = AppSettings()
+        assert hasattr(app_settings, 'log_level'), \
+            "AppSettings should have log_level field"
+
+        # Verify otel_config reads log level
+        config = OpenTelemetryConfig()
+        assert hasattr(config, 'log_level'), \
+            "OpenTelemetryConfig should have log_level attribute"
+        assert isinstance(config.log_level, int), \
+            "log_level should be an integer (logging level)"
+
+        # Verify log level is one of the valid logging levels
+        valid_levels = [logging.DEBUG, logging.INFO, logging.WARNING, logging.ERROR, logging.CRITICAL]
+        assert config.log_level in valid_levels, \
+            f"log_level should be a valid logging level, got {config.log_level}"

atlas/tests/test_mcp_auth_routes.py

@@ -0,0 +1,341 @@
+"""Unit tests for MCP authentication routes.
+
+Tests the API endpoints for per-user token management:
+- GET /api/mcp/auth/status - Get auth status for all servers
+- POST /api/mcp/auth/{server}/token - Upload token for server
+- DELETE /api/mcp/auth/{server}/token - Remove token for server
+
+Updated: 2025-01-21
+"""
+
+import time
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from atlas.core.log_sanitizer import get_current_user
+from atlas.routes.mcp_auth_routes import TokenUpload, router
+
+
+# Create a test app with the auth routes
+def create_test_app(user_override: str = "test@example.com"):
+    """Create a FastAPI test app with auth routes."""
+    app = FastAPI()
+    app.include_router(router)
+
+    # Override the get_current_user dependency
+    async def override_get_current_user():
+        return user_override
+
+    app.dependency_overrides[get_current_user] = override_get_current_user
+    return app
+
+
+class TestGetAuthStatus:
+    """Test GET /api/mcp/auth/status endpoint."""
+
+    @pytest.fixture
+    def client(self):
+        """Create test client."""
+        app = create_test_app()
+        return TestClient(app)
+
+    @pytest.fixture
+    def mock_dependencies(self):
+        """Mock the dependencies for auth routes."""
+        with patch("atlas.routes.mcp_auth_routes.app_factory") as mock_factory, \
+             patch("atlas.routes.mcp_auth_routes.get_token_storage") as mock_storage:
+
+            # Mock MCP manager
+            mock_mcp_manager = AsyncMock()
+            mock_mcp_manager.get_authorized_servers = AsyncMock(return_value=["server1", "server2"])
+            mock_mcp_manager.servers_config = {
+                "server1": {"auth_type": "api_key", "description": "API Key Server"},
+                "server2": {"auth_type": "jwt", "description": "JWT Server"},
+            }
+            mock_factory.get_mcp_manager.return_value = mock_mcp_manager
+
+            # Mock token storage
+            mock_token_storage = MagicMock()
+            mock_token_storage.get_user_auth_status.return_value = {
+                "server1": {
+                    "token_type": "api_key",
+                    "is_expired": False,
+                    "expires_at": None,
+                    "time_until_expiry": None,
+                    "has_refresh_token": False,
+                    "scopes": None,
+                }
+            }
+            mock_storage.return_value = mock_token_storage
+
+            yield {
+                "factory": mock_factory,
+                "storage": mock_storage,
+                "mcp_manager": mock_mcp_manager,
+                "token_storage": mock_token_storage,
+            }
+
+    def test_get_auth_status_success(self, client, mock_dependencies):
+        """Should return auth status for all servers."""
+        response = client.get("/api/mcp/auth/status")
+
+        assert response.status_code == 200
+        data = response.json()
+
+        assert "servers" in data
+        assert "user" in data
+        assert data["user"] == "test@example.com"
+        assert len(data["servers"]) == 2
+
+    def test_get_auth_status_shows_authenticated_servers(self, client, mock_dependencies):
+        """Should indicate which servers user is authenticated with."""
+        response = client.get("/api/mcp/auth/status")
+
+        data = response.json()
+        servers = {s["server_name"]: s for s in data["servers"]}
+
+        # server1 has token stored
+        assert servers["server1"]["authenticated"] is True
+        assert servers["server1"]["auth_type"] == "api_key"
+
+        # server2 has no token
+        assert servers["server2"]["authenticated"] is False
+        assert servers["server2"]["auth_type"] == "jwt"
+
+    def test_get_auth_status_includes_token_details(self, client, mock_dependencies):
+        """Should include token details for authenticated servers."""
+        response = client.get("/api/mcp/auth/status")
+
+        data = response.json()
+        server1 = next(s for s in data["servers"] if s["server_name"] == "server1")
+
+        assert server1["token_type"] == "api_key"
+        assert server1["is_expired"] is False
+
+
+class TestUploadToken:
+    """Test POST /api/mcp/auth/{server_name}/token endpoint."""
+
+    @pytest.fixture
+    def client(self):
+        """Create test client."""
+        app = create_test_app()
+        return TestClient(app)
+
+    @pytest.fixture
+    def mock_dependencies(self):
+        """Mock the dependencies for auth routes."""
+        with patch("atlas.routes.mcp_auth_routes.app_factory") as mock_factory, \
+             patch("atlas.routes.mcp_auth_routes.get_token_storage") as mock_storage:
+
+            mock_mcp_manager = AsyncMock()
+            mock_mcp_manager.get_authorized_servers = AsyncMock(return_value=["test-server"])
+            mock_mcp_manager.servers_config = {
+                "test-server": {"auth_type": "api_key", "description": "Test Server"},
+            }
+            mock_factory.get_mcp_manager.return_value = mock_mcp_manager
+
+            mock_token_storage = MagicMock()
+            mock_stored_token = MagicMock()
+            mock_stored_token.token_type = "api_key"
+            mock_stored_token.expires_at = None
+            mock_stored_token.scopes = None
+            mock_token_storage.store_token.return_value = mock_stored_token
+            mock_storage.return_value = mock_token_storage
+
+            yield {
+                "factory": mock_factory,
+                "storage": mock_storage,
+                "mcp_manager": mock_mcp_manager,
+                "token_storage": mock_token_storage,
+            }
+
+    def test_upload_token_success(self, client, mock_dependencies):
+        """Should store token successfully."""
+        response = client.post(
+            "/api/mcp/auth/test-server/token",
+            json={"token": "my-api-key-123"}
+        )
+
+        assert response.status_code == 200
+        data = response.json()
+
+        assert data["message"] == "Token stored for server 'test-server'"
+        assert data["server_name"] == "test-server"
+        assert data["token_type"] == "api_key"
+
+    def test_upload_token_with_expiry(self, client, mock_dependencies):
+        """Should store token with expiration time."""
+        expiry = time.time() + 3600
+        response = client.post(
+            "/api/mcp/auth/test-server/token",
+            json={"token": "my-api-key", "expires_at": expiry}
+        )
+
+        assert response.status_code == 200
+        mock_dependencies["token_storage"].store_token.assert_called_once()
+        call_args = mock_dependencies["token_storage"].store_token.call_args
+        assert call_args.kwargs["expires_at"] == expiry
+
+    def test_upload_token_with_scopes(self, client, mock_dependencies):
+        """Should store token with scopes."""
+        response = client.post(
+            "/api/mcp/auth/test-server/token",
+            json={"token": "my-api-key", "scopes": "read write"}
+        )
+
+        assert response.status_code == 200
+        mock_dependencies["token_storage"].store_token.assert_called_once()
+        call_args = mock_dependencies["token_storage"].store_token.call_args
+        assert call_args.kwargs["scopes"] == "read write"
+
+    def test_upload_token_empty_rejected(self, client, mock_dependencies):
+        """Should reject empty token."""
+        response = client.post(
+            "/api/mcp/auth/test-server/token",
+            json={"token": ""}
+        )
+
+        assert response.status_code == 400
+        assert "empty" in response.json()["detail"].lower()
+
+    def test_upload_token_whitespace_only_rejected(self, client, mock_dependencies):
+        """Should reject whitespace-only token."""
+        response = client.post(
+            "/api/mcp/auth/test-server/token",
+            json={"token": "   "}
+        )
+
+        assert response.status_code == 400
+
+    def test_upload_token_unauthorized_server(self, client, mock_dependencies):
+        """Should reject token for unauthorized server."""
+        response = client.post(
+            "/api/mcp/auth/unauthorized-server/token",
+            json={"token": "my-api-key"}
+        )
+
+        assert response.status_code == 403
+        assert "Not authorized" in response.json()["detail"]
+
+    def test_upload_token_wrong_auth_type(self, client, mock_dependencies):
+        """Should reject token for server with auth_type=none."""
+        mock_dependencies["mcp_manager"].servers_config["test-server"]["auth_type"] = "none"
+
+        response = client.post(
+            "/api/mcp/auth/test-server/token",
+            json={"token": "my-api-key"}
+        )
+
+        assert response.status_code == 400
+        assert "does not accept token authentication" in response.json()["detail"]
+
+    def test_upload_token_strips_whitespace(self, client, mock_dependencies):
+        """Should strip whitespace from token."""
+        response = client.post(
+            "/api/mcp/auth/test-server/token",
+            json={"token": "  my-api-key  "}
+        )
+
+        assert response.status_code == 200
+        call_args = mock_dependencies["token_storage"].store_token.call_args
+        assert call_args.kwargs["token_value"] == "my-api-key"
+
+
+class TestRemoveToken:
+    """Test DELETE /api/mcp/auth/{server_name}/token endpoint."""
+
+    @pytest.fixture
+    def client(self):
+        """Create test client."""
+        app = create_test_app()
+        return TestClient(app)
+
+    @pytest.fixture
+    def mock_dependencies(self):
+        """Mock the dependencies for auth routes."""
+        with patch("atlas.routes.mcp_auth_routes.get_token_storage") as mock_storage, \
+             patch("atlas.routes.mcp_auth_routes.app_factory") as mock_factory:
+
+            mock_token_storage = MagicMock()
+            mock_token_storage.remove_token.return_value = True
+            mock_storage.return_value = mock_token_storage
+
+            # Mock tool manager for cache invalidation
+            mock_tool_manager = AsyncMock()
+            mock_tool_manager._invalidate_user_client = AsyncMock()
+            mock_factory.get_mcp_manager.return_value = mock_tool_manager
+
+            yield {
+                "storage": mock_storage,
+                "token_storage": mock_token_storage,
+                "factory": mock_factory,
+                "tool_manager": mock_tool_manager,
+            }
+
+    def test_remove_token_success(self, client, mock_dependencies):
+        """Should remove token successfully."""
+        response = client.delete("/api/mcp/auth/test-server/token")
+
+        assert response.status_code == 200
+        data = response.json()
+
+        assert data["message"] == "Token removed for server 'test-server'"
+        assert data["server_name"] == "test-server"
+
+    def test_remove_token_invalidates_cache(self, client, mock_dependencies):
+        """Should invalidate cached client when token is removed."""
+        response = client.delete("/api/mcp/auth/test-server/token")
+
+        assert response.status_code == 200
+        # Verify cache invalidation was called
+        mock_dependencies["tool_manager"]._invalidate_user_client.assert_called_once_with(
+            "test@example.com", "test-server"
+        )
+
+    def test_remove_token_not_found(self, client, mock_dependencies):
+        """Should return 404 when token doesn't exist."""
+        mock_dependencies["token_storage"].remove_token.return_value = False
+
+        response = client.delete("/api/mcp/auth/nonexistent-server/token")
+
+        assert response.status_code == 404
+        assert "No token found" in response.json()["detail"]
+
+
+class TestTokenUploadModel:
+    """Test TokenUpload Pydantic model."""
+
+    def test_token_required(self):
+        """Token field should be required."""
+        with pytest.raises(Exception):
+            TokenUpload()
+
+    def test_token_accepts_string(self):
+        """Token field should accept string."""
+        model = TokenUpload(token="my-api-key")
+        assert model.token == "my-api-key"
+
+    def test_expires_at_optional(self):
+        """expires_at should be optional."""
+        model = TokenUpload(token="my-api-key")
+        assert model.expires_at is None
+
+    def test_expires_at_accepts_float(self):
+        """expires_at should accept float timestamp."""
+        expiry = time.time() + 3600
+        model = TokenUpload(token="my-api-key", expires_at=expiry)
+        assert model.expires_at == expiry
+
+    def test_scopes_optional(self):
+        """scopes should be optional."""
+        model = TokenUpload(token="my-api-key")
+        assert model.scopes is None
+
+    def test_scopes_accepts_string(self):
+        """scopes should accept space-separated string."""
+        model = TokenUpload(token="my-api-key", scopes="read write admin")
+        assert model.scopes == "read write admin"