atlas-chat 0.1.0__py3-none-any.whl

atlas/tests/test_backend_public_url.py

@@ -0,0 +1,185 @@
+"""
+Tests for BACKEND_PUBLIC_URL configuration and absolute URL generation.
+
+This module tests that the file download URL generation correctly handles:
+1. Relative URLs when BACKEND_PUBLIC_URL is not configured (backward compatibility)
+2. Absolute URLs when BACKEND_PUBLIC_URL is configured (remote MCP server support)
+3. Proper token generation and validation
+"""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from atlas.core.capabilities import create_download_url, generate_file_token
+
+
+def create_mock_settings(backend_public_url=None):
+    """Helper to create a properly mocked settings object."""
+    mock_settings = MagicMock()
+    mock_settings.backend_public_url = backend_public_url
+    mock_settings.capability_token_secret = "test-secret-key-for-testing"
+    mock_settings.capability_token_ttl_seconds = 3600
+    return mock_settings
+
+
+class TestBackendPublicUrlConfiguration:
+    """Test suite for BACKEND_PUBLIC_URL configuration behavior."""
+
+    def test_relative_url_without_backend_public_url(self):
+        """Test that relative URLs are generated when BACKEND_PUBLIC_URL is not set."""
+        mock_settings = create_mock_settings(backend_public_url=None)
+
+        with patch('atlas.core.capabilities.config_manager') as mock_cm:
+            mock_cm.app_settings = mock_settings
+
+            url = create_download_url("test-key-123", "user@example.com")
+
+            # Should start with / (relative URL)
+            assert url.startswith("/api/files/download/")
+            assert "test-key-123" in url
+            assert "token=" in url
+            assert not url.startswith("http")
+
+    def test_absolute_url_with_backend_public_url(self):
+        """Test that absolute URLs are generated when BACKEND_PUBLIC_URL is configured."""
+        mock_settings = create_mock_settings(backend_public_url="https://atlas.example.com")
+
+        with patch('atlas.core.capabilities.config_manager') as mock_cm:
+            mock_cm.app_settings = mock_settings
+
+            url = create_download_url("test-key-456", "admin@example.com")
+
+            # Should be absolute URL
+            assert url.startswith("https://atlas.example.com/api/files/download/")
+            assert "test-key-456" in url
+            assert "token=" in url
+
+    def test_absolute_url_strips_trailing_slash(self):
+        """Test that trailing slashes in BACKEND_PUBLIC_URL are handled correctly."""
+        mock_settings = create_mock_settings(backend_public_url="https://atlas.example.com/")
+
+        with patch('atlas.core.capabilities.config_manager') as mock_cm:
+            mock_cm.app_settings = mock_settings
+
+            url = create_download_url("test-key-789", "user@example.com")
+
+            # Should not have double slash
+            assert "https://atlas.example.com/api/files/download/" in url
+            assert "https://atlas.example.com//api" not in url
+
+    def test_url_with_non_standard_port(self):
+        """Test absolute URL generation with non-standard port."""
+        mock_settings = create_mock_settings(backend_public_url="https://atlas.example.com:8443")
+
+        with patch('atlas.core.capabilities.config_manager') as mock_cm:
+            mock_cm.app_settings = mock_settings
+
+            url = create_download_url("test-key-abc", "user@example.com")
+
+            # Should include port in URL
+            assert url.startswith("https://atlas.example.com:8443/api/files/download/")
+
+    def test_url_with_localhost(self):
+        """Test absolute URL generation with localhost (development mode)."""
+        mock_settings = create_mock_settings(backend_public_url="http://localhost:8000")
+
+        with patch('atlas.core.capabilities.config_manager') as mock_cm:
+            mock_cm.app_settings = mock_settings
+
+            url = create_download_url("test-key-dev", "dev@example.com")
+
+            # Should use localhost URL
+            assert url.startswith("http://localhost:8000/api/files/download/")
+
+    def test_fallback_without_user_email(self):
+        """Test URL generation without user email (no token)."""
+        mock_settings = create_mock_settings(backend_public_url="https://atlas.example.com")
+
+        with patch('atlas.core.capabilities.config_manager') as mock_cm:
+            mock_cm.app_settings = mock_settings
+
+            url = create_download_url("test-key-nouser", None)
+
+            # Should be absolute but without token
+            assert url.startswith("https://atlas.example.com/api/files/download/")
+            assert "test-key-nouser" in url
+            assert "token=" not in url
+
+    def test_config_error_falls_back_to_relative(self):
+        """Test that configuration errors fall back to relative URLs gracefully."""
+        # Mock config manager app_settings to raise exception when accessed
+        mock_cm = MagicMock()
+        # When app_settings is accessed, raise exception
+        type(mock_cm).app_settings = property(lambda self: (_ for _ in ()).throw(Exception("Config error")))
+
+        # Also need to mock _get_secret since it also accesses config_manager
+        with patch('atlas.core.capabilities.config_manager', mock_cm):
+            with patch('atlas.core.capabilities._get_secret', return_value=b'test-secret'):
+                url = create_download_url("test-key-error", "user@example.com")
+
+                # Should fall back to relative URL
+                assert url.startswith("/api/files/download/")
+                assert not url.startswith("http")
+
+
+class TestTokenValidation:
+    """Test suite for token generation and validation with absolute URLs."""
+
+    def test_token_works_with_absolute_urls(self):
+        """Test that tokens generated for absolute URLs are valid."""
+        from atlas.core.capabilities import verify_file_token
+
+        user_email = "test@example.com"
+        file_key = "test-key-123"
+
+        mock_settings = create_mock_settings()
+        with patch('atlas.core.capabilities.config_manager') as mock_cm:
+            mock_cm.app_settings = mock_settings
+
+            # Generate token
+            token = generate_file_token(user_email, file_key, ttl_seconds=60)
+
+            # Verify token
+            claims = verify_file_token(token)
+
+            assert claims is not None
+            assert claims["u"] == user_email
+            assert claims["k"] == file_key
+            assert claims["e"] > 0  # Expiry timestamp exists
+
+
+class TestBackwardCompatibility:
+    """Test suite to ensure backward compatibility with existing behavior."""
+
+    def test_stdio_servers_still_work_with_relative_urls(self):
+        """Test that stdio (local) servers continue to work with relative URLs."""
+        mock_settings = create_mock_settings(backend_public_url=None)
+
+        with patch('atlas.core.capabilities.config_manager') as mock_cm:
+            mock_cm.app_settings = mock_settings
+
+            url = create_download_url("local-file", "local@example.com")
+
+            # Local servers can resolve relative URLs
+            assert url.startswith("/api/files/download/")
+
+    def test_existing_mcp_servers_unaffected(self):
+        """Test that existing MCP server configurations continue to work."""
+        # This test verifies that the changes don't break existing deployments
+        # that haven't configured BACKEND_PUBLIC_URL
+
+        mock_settings = create_mock_settings(backend_public_url=None)
+
+        with patch('atlas.core.capabilities.config_manager') as mock_cm:
+            mock_cm.app_settings = mock_settings
+
+            # Should handle missing attribute gracefully and return relative URL
+            url = create_download_url("legacy-key", "legacy@example.com")
+
+            assert url.startswith("/api/files/download/")
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])
+

atlas/tests/test_banner_logging.py

@@ -0,0 +1,287 @@
+"""Tests for banner message save logging functionality."""
+import logging
+
+from main import app
+from starlette.testclient import TestClient
+
+from atlas.modules.config import config_manager
+
+
+def test_banner_save_success_logging(caplog, tmp_path, monkeypatch):
+    """Test that successful banner save produces INFO log with file path."""
+    client = TestClient(app)
+
+    # Setup temp config directory
+    config_dir = tmp_path / "config" / "overrides"
+    config_dir.mkdir(parents=True, exist_ok=True)
+
+    # Mock config path to use temp directory
+    def mock_get_admin_config_path(filename):
+        return config_dir / filename
+
+    monkeypatch.setattr(
+        "atlas.routes.admin_routes.get_admin_config_path",
+        mock_get_admin_config_path
+    )
+
+    # Mock setup_config_overrides to avoid side effects
+    monkeypatch.setattr("atlas.routes.admin_routes.setup_config_overrides", lambda: None)
+
+    # Capture logs at INFO level
+    with caplog.at_level(logging.INFO):
+        # Make request to update banner messages
+        response = client.post(
+            "/admin/banners",
+            json={"messages": ["Test banner message", "Another test message"]},
+            headers={"X-User-Email": "admin@example.com"}
+        )
+
+    # Verify request succeeded
+    assert response.status_code == 200
+
+    # Check that INFO log was created with success message
+    info_logs = [record for record in caplog.records if record.levelname == "INFO"]
+    assert len(info_logs) > 0
+
+    # Find the banner save log
+    banner_logs = [
+        log for log in info_logs
+        if "Banner messages successfully saved to disk" in log.message
+    ]
+    assert len(banner_logs) == 1
+
+    # Verify log contains file path and admin user
+    log_message = banner_logs[0].message
+    assert "messages.txt" in log_message
+    assert "admin@example.com" in log_message
+
+
+def test_banner_save_failure_logging(caplog, tmp_path, monkeypatch):
+    """Test that failed banner save produces ERROR log with details."""
+    client = TestClient(app)
+
+    # Mock get_admin_config_path to return a path
+    readonly_file = tmp_path / "readonly.txt"
+    readonly_file.write_text("test")
+
+    def mock_get_admin_config_path(filename):
+        return readonly_file
+
+    monkeypatch.setattr(
+        "atlas.routes.admin_routes.get_admin_config_path",
+        mock_get_admin_config_path
+    )
+
+    # Mock setup_config_overrides to avoid side effects
+    monkeypatch.setattr("atlas.routes.admin_routes.setup_config_overrides", lambda: None)
+
+    # Mock write_file_content to raise an exception
+    def mock_write_file_content(file_path, content, file_type="text"):
+        raise PermissionError("Permission denied")
+
+    monkeypatch.setattr(
+        "atlas.routes.admin_routes.write_file_content",
+        mock_write_file_content
+    )
+
+    # Capture logs at ERROR level
+    with caplog.at_level(logging.ERROR):
+        # Make request to update banner messages (should fail)
+        response = client.post(
+            "/admin/banners",
+            json={"messages": ["Test banner message"]},
+            headers={"X-User-Email": "admin@example.com"}
+        )
+
+    # Verify request failed
+    assert response.status_code == 500
+
+    # Check that ERROR log was created with failure message
+    error_logs = [record for record in caplog.records if record.levelname == "ERROR"]
+    assert len(error_logs) > 0
+
+    # Find the banner save error log
+    banner_error_logs = [
+        log for log in error_logs
+        if "Failed to save banner messages to disk" in log.message
+    ]
+    assert len(banner_error_logs) == 1
+
+    # Verify log contains file path and error details
+    log_message = banner_error_logs[0].message
+    assert "readonly.txt" in log_message
+    assert "Permission denied" in log_message or "PermissionError" in log_message
+
+
+def test_banner_save_logs_sanitized_paths(caplog, tmp_path, monkeypatch):
+    """Test that file paths in logs are sanitized to prevent log injection."""
+    client = TestClient(app)
+
+    # Setup temp config directory with a potentially malicious name
+    config_dir = tmp_path / "config" / "overrides"
+    config_dir.mkdir(parents=True, exist_ok=True)
+
+    # Mock config path
+    def mock_get_admin_config_path(filename):
+        return config_dir / filename
+
+    monkeypatch.setattr(
+        "atlas.routes.admin_routes.get_admin_config_path",
+        mock_get_admin_config_path
+    )
+
+    # Mock setup_config_overrides to avoid side effects
+    monkeypatch.setattr("atlas.routes.admin_routes.setup_config_overrides", lambda: None)
+
+    # Capture logs at INFO level
+    with caplog.at_level(logging.INFO):
+        # Make request to update banner messages
+        response = client.post(
+            "/admin/banners",
+            json={"messages": ["Test message"]},
+            headers={"X-User-Email": "admin@example.com"}
+        )
+
+    # Verify request succeeded
+    assert response.status_code == 200
+
+    # Check that log messages don't contain raw newlines or control characters
+    info_logs = [record for record in caplog.records if record.levelname == "INFO"]
+    banner_logs = [
+        log for log in info_logs
+        if "Banner messages successfully saved to disk" in log.message
+    ]
+
+    assert len(banner_logs) == 1
+    log_message = banner_logs[0].message
+
+    # Verify no newlines in the log message
+    assert "\n" not in log_message
+    assert "\r" not in log_message
+
+
+def test_banner_get_includes_enabled_status(tmp_path, monkeypatch):
+    """Test that GET /admin/banners includes banner_enabled status."""
+    client = TestClient(app)
+
+    # Setup temp config directory
+    config_dir = tmp_path / "config" / "overrides"
+    config_dir.mkdir(parents=True, exist_ok=True)
+    messages_file = config_dir / "messages.txt"
+    messages_file.write_text("Test message\n")
+
+    # Mock config path to use temp directory
+    def mock_get_admin_config_path(filename):
+        return config_dir / filename
+
+    monkeypatch.setattr(
+        "atlas.routes.admin_routes.get_admin_config_path",
+        mock_get_admin_config_path
+    )
+
+    # Mock setup_config_overrides to avoid side effects
+    monkeypatch.setattr("atlas.routes.admin_routes.setup_config_overrides", lambda: None)
+
+    # Make request to get banner config
+    response = client.get(
+        "/admin/banners",
+        headers={"X-User-Email": "admin@example.com"}
+    )
+
+    # Verify request succeeded
+    assert response.status_code == 200
+
+    # Check response contains banner_enabled field
+    data = response.json()
+    assert "banner_enabled" in data
+    assert isinstance(data["banner_enabled"], bool)
+    # The field should match the current config setting
+    assert data["banner_enabled"] == config_manager.app_settings.banner_enabled
+
+
+def test_banner_get_with_feature_disabled(tmp_path, monkeypatch):
+    """Test that GET /admin/banners returns banner_enabled: false when feature is disabled."""
+    client = TestClient(app)
+
+    # Setup temp config directory
+    config_dir = tmp_path / "config" / "overrides"
+    config_dir.mkdir(parents=True, exist_ok=True)
+    messages_file = config_dir / "messages.txt"
+    messages_file.write_text("Test message\n")
+
+    # Mock config path
+    def mock_get_admin_config_path(filename):
+        return config_dir / filename
+
+    monkeypatch.setattr(
+        "atlas.routes.admin_routes.get_admin_config_path",
+        mock_get_admin_config_path
+    )
+
+    # Mock setup_config_overrides
+    monkeypatch.setattr("atlas.routes.admin_routes.setup_config_overrides", lambda: None)
+
+    # Mock banner_enabled to be false
+    monkeypatch.setattr(
+        "atlas.routes.admin_routes.config_manager.app_settings.banner_enabled",
+        False
+    )
+
+    # Make request to get banner config
+    response = client.get(
+        "/admin/banners",
+        headers={"X-User-Email": "admin@example.com"}
+    )
+
+    # Verify request succeeded
+    assert response.status_code == 200
+
+    # Check response contains banner_enabled field set to false
+    data = response.json()
+    assert "banner_enabled" in data
+    assert data["banner_enabled"] is False
+
+
+def test_banner_get_with_feature_enabled(tmp_path, monkeypatch):
+    """Test that GET /admin/banners returns banner_enabled: true when feature is enabled."""
+    client = TestClient(app)
+
+    # Setup temp config directory
+    config_dir = tmp_path / "config" / "overrides"
+    config_dir.mkdir(parents=True, exist_ok=True)
+    messages_file = config_dir / "messages.txt"
+    messages_file.write_text("Test message\n")
+
+    # Mock config path
+    def mock_get_admin_config_path(filename):
+        return config_dir / filename
+
+    monkeypatch.setattr(
+        "atlas.routes.admin_routes.get_admin_config_path",
+        mock_get_admin_config_path
+    )
+
+    # Mock setup_config_overrides
+    monkeypatch.setattr("atlas.routes.admin_routes.setup_config_overrides", lambda: None)
+
+    # Mock banner_enabled to be true
+    monkeypatch.setattr(
+        "atlas.routes.admin_routes.config_manager.app_settings.banner_enabled",
+        True
+    )
+
+    # Make request to get banner config
+    response = client.get(
+        "/admin/banners",
+        headers={"X-User-Email": "admin@example.com"}
+    )
+
+    # Verify request succeeded
+    assert response.status_code == 200
+
+    # Check response contains banner_enabled field set to true
+    data = response.json()
+    assert "banner_enabled" in data
+    assert data["banner_enabled"] is True
+
+

atlas/tests/test_capability_tokens_and_injection.py

@@ -0,0 +1,203 @@
+import base64
+import os
+import sys
+import types
+from typing import Optional
+
+import pytest
+from fastapi.testclient import TestClient
+
+# Ensure backend root is on path
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+# Save original module before patching (so we can restore it later)
+_ORIGINAL_LITELLM_MODULE = sys.modules.get("atlas.modules.llm.litellm_caller")
+
+# Stub LiteLLM before importing app to avoid external dependency in tests
+fake_litellm_caller = types.ModuleType("atlas.modules.llm.litellm_caller")
+
+
+class _FakeLLM:
+    def __init__(self, *args, **kwargs):
+        pass
+
+    async def call_plain(self, *args, **kwargs):
+        return "ok"
+
+
+fake_litellm_caller.LiteLLMCaller = _FakeLLM  # type: ignore
+sys.modules["atlas.modules.llm.litellm_caller"] = fake_litellm_caller
+
+from main import app  # noqa: E402  # type: ignore
+
+
+def _restore_original_module():
+    """Restore the original litellm_caller module if it was saved."""
+    if _ORIGINAL_LITELLM_MODULE is not None:
+        sys.modules["atlas.modules.llm.litellm_caller"] = _ORIGINAL_LITELLM_MODULE
+    elif "atlas.modules.llm.litellm_caller" in sys.modules:
+        del sys.modules["atlas.modules.llm.litellm_caller"]
+
+
+@pytest.fixture(scope="module", autouse=True)
+def restore_litellm_module_after_tests():
+    """Restore the original LiteLLM module after all tests in this module complete."""
+    yield
+    _restore_original_module()
+
+from atlas.core.capabilities import generate_file_token, verify_file_token  # noqa: E402  # type: ignore
+
+
+class FakeS3:
+    def __init__(self):
+        self._store = {}
+        self.endpoint_url = "mock://s3"
+        self.bucket_name = "test-bucket"
+
+    async def upload_file(self, user_email: str, filename: str, content_base64: str, content_type: str = "application/octet-stream", tags=None, source_type: str = "user"):
+        key = f"k_{len(self._store)+1}"
+        meta = {
+            "key": key,
+            "filename": filename,
+            "size": len(base64.b64decode(content_base64)),
+            "content_type": content_type,
+            "last_modified": "now",
+            "etag": "test",
+            "tags": tags or {"source": source_type},
+            "user_email": user_email,
+            "content_base64": content_base64,
+        }
+        self._store[key] = meta
+        return meta
+
+    async def get_file(self, user_email: str, file_key: str):
+        return self._store.get(file_key)
+
+
+@pytest.fixture()
+def client(monkeypatch):
+    # Inject fake S3 into app_factory for route handlers
+    from atlas.infrastructure import app_factory as af  # type: ignore
+
+    fake = FakeS3()
+
+    original_get = af.get_file_storage
+    af.get_file_storage = lambda: fake  # type: ignore
+
+    try:
+        yield TestClient(app)
+    finally:
+        # Restore
+        af.get_file_storage = original_get  # type: ignore
+
+
+def test_capability_token_roundtrip():
+    token = generate_file_token("alice@example.com", "k123", ttl_seconds=60)
+    claims = verify_file_token(token)
+    assert claims is not None
+    assert claims["u"] == "alice@example.com"
+    assert claims["k"] == "k123"
+
+
+def _extract_key(resp_json) -> Optional[str]:
+    if isinstance(resp_json, dict):
+        return resp_json.get("key")
+    return None
+
+
+def test_download_with_and_without_token(client):
+    """Test file upload/download with proper authentication in production mode."""
+    # Upload a small text file via API with authentication header
+    content = base64.b64encode(b"hello world").decode("utf-8")
+    upload_resp = client.post(
+        "/api/files",
+        json={
+            "filename": "hello.txt",
+            "content_base64": content,
+            "content_type": "text/plain",
+            "tags": {"source": "test"},
+        },
+        headers={"X-User-Email": "test@example.com"},
+    )
+    assert upload_resp.status_code == 200
+    key = _extract_key(upload_resp.json())
+    assert key
+
+    # Download with same user authentication
+    dl_resp = client.get(
+        f"/api/files/download/{key}",
+        headers={"X-User-Email": "test@example.com"}
+    )
+    assert dl_resp.status_code == 200
+    assert dl_resp.content == b"hello world"
+
+    # With capability token for a different user (token grants access)
+    token = generate_file_token("alice@example.com", key, ttl_seconds=60)
+    dl_resp2 = client.get(
+        f"/api/files/download/{key}",
+        params={"token": token},
+        headers={"X-User-Email": "alice@example.com"}
+    )
+    assert dl_resp2.status_code == 200
+    assert dl_resp2.content == b"hello world"
+
+
+def test_download_without_auth_fails_in_production_mode(client):
+    """Test that requests without authentication fail in production mode (debug_mode=False).
+
+    This test validates that when debug_mode=False (production mode), requests without
+    the X-User-Email header are rejected with 401 Unauthorized.
+
+    Note: This test will pass in production mode and fail in debug mode, which is expected.
+    CI runs tests in both modes to validate both behaviors.
+    """
+    from atlas.infrastructure.app_factory import app_factory  # type: ignore
+
+    # Skip this test if running in debug mode (it's expected to fail)
+    if app_factory.config_manager.app_settings.debug_mode:
+        pytest.skip("Skipping production mode test - running in debug mode")
+
+    content = base64.b64encode(b"test content").decode("utf-8")
+
+    # Try to upload without authentication - should fail with 401 in production mode
+    upload_resp = client.post(
+        "/api/files",
+        json={
+            "filename": "test.txt",
+            "content_base64": content,
+            "content_type": "text/plain",
+            "tags": {"source": "test"},
+        },
+    )
+    assert upload_resp.status_code == 401
+
+
+def test_injection_produces_tokenized_urls(client, monkeypatch):
+    """Verify that tool argument injection replaces filename with tokenized URL."""
+    from atlas.application.chat.utilities.tool_executor import inject_context_into_args  # type: ignore
+
+    # Create a fake session context with a file mapping
+    session_context = {
+        "user_email": "bob@example.com",
+        "files": {
+            "report.pdf": {"key": "abc123", "content_type": "application/pdf"}
+        },
+    }
+
+    args = {"filename": "report.pdf"}
+    injected = inject_context_into_args(args, session_context)
+
+    assert injected["username"] == "bob@example.com"
+    assert injected["original_filename"] == "report.pdf"
+    # URL should include /api/files/download/abc123 and a token query param
+    assert injected["filename"].startswith("/api/files/download/abc123")
+    assert "?token=" in injected["filename"]
+
+    # Multiple files
+    args2 = {"file_names": ["report.pdf", "missing.txt"]}
+    injected2 = inject_context_into_args(args2, session_context)
+    assert injected2["original_file_names"] == ["report.pdf", "missing.txt"]
+    assert injected2["file_names"][0].startswith("/api/files/download/abc123")
+    assert "?token=" in injected2["file_names"][0]
+    # Missing.txt doesn't resolve to key, kept as-is
+    assert injected2["file_names"][1] == "missing.txt"