feat: use XDG directory standards

Signed-off-by: Mustafa Elbehery <melbeher@redhat.com>

tests/integration/test_xdg_e2e.py

@@ -0,0 +1,593 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import os
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+import yaml
+
+from llama_stack.distribution.utils.xdg_utils import (
+    get_llama_stack_config_dir,
+    get_llama_stack_data_dir,
+    get_llama_stack_state_dir,
+)
+
+
+class TestXDGEndToEnd(unittest.TestCase):
+    """End-to-end tests for XDG compliance workflows."""
+
+    def setUp(self):
+        """Set up test environment."""
+        self.original_env = {}
+        self.env_vars = [
+            "XDG_CONFIG_HOME",
+            "XDG_DATA_HOME",
+            "XDG_STATE_HOME",
+            "XDG_CACHE_HOME",
+            "LLAMA_STACK_CONFIG_DIR",
+            "SQLITE_STORE_DIR",
+            "FILES_STORAGE_DIR",
+        ]
+
+        for key in self.env_vars:
+            self.original_env[key] = os.environ.get(key)
+
+    def tearDown(self):
+        """Clean up test environment."""
+        for key, value in self.original_env.items():
+            if value is None:
+                os.environ.pop(key, None)
+            else:
+                os.environ[key] = value
+
+    def clear_env_vars(self):
+        """Clear all relevant environment variables."""
+        for key in self.env_vars:
+            os.environ.pop(key, None)
+
+    def create_realistic_legacy_structure(self, base_dir: Path) -> Path:
+        """Create a realistic legacy ~/.llama directory structure."""
+        legacy_dir = base_dir / ".llama"
+        legacy_dir.mkdir()
+
+        # Create distributions with realistic content
+        distributions_dir = legacy_dir / "distributions"
+        distributions_dir.mkdir()
+
+        # Ollama distribution
+        ollama_dir = distributions_dir / "ollama"
+        ollama_dir.mkdir()
+
+        ollama_run_yaml = ollama_dir / "ollama-run.yaml"
+        ollama_run_yaml.write_text("""
+version: 2
+apis:
+  - inference
+  - safety
+  - memory
+  - vector_io
+  - agents
+  - files
+providers:
+  inference:
+    - provider_type: remote::ollama
+      config:
+        url: http://localhost:11434
+""")
+
+        ollama_build_yaml = ollama_dir / "build.yaml"
+        ollama_build_yaml.write_text("""
+name: ollama
+description: Ollama inference provider
+docker_image: ollama:latest
+""")
+
+        # Create providers.d structure
+        providers_dir = legacy_dir / "providers.d"
+        providers_dir.mkdir()
+
+        remote_dir = providers_dir / "remote"
+        remote_dir.mkdir()
+
+        inference_dir = remote_dir / "inference"
+        inference_dir.mkdir()
+
+        custom_provider = inference_dir / "custom-inference.yaml"
+        custom_provider.write_text("""
+provider_type: remote::custom
+config:
+  url: http://localhost:8080
+  api_key: test_key
+""")
+
+        # Create checkpoints with model files
+        checkpoints_dir = legacy_dir / "checkpoints"
+        checkpoints_dir.mkdir()
+
+        model_dir = checkpoints_dir / "meta-llama" / "Llama-3.2-1B-Instruct"
+        model_dir.mkdir(parents=True)
+
+        # Create fake model files
+        (model_dir / "consolidated.00.pth").write_bytes(b"fake model weights" * 1000)
+        (model_dir / "params.json").write_text('{"dim": 2048, "n_layers": 22}')
+        (model_dir / "tokenizer.model").write_bytes(b"fake tokenizer" * 100)
+
+        # Create runtime with databases
+        runtime_dir = legacy_dir / "runtime"
+        runtime_dir.mkdir()
+
+        (runtime_dir / "trace_store.db").write_text("SQLite format 3\x00" + "fake database content")
+        (runtime_dir / "agent_sessions.db").write_text("SQLite format 3\x00" + "fake agent sessions")
+
+        # Create config files
+        (legacy_dir / "config.json").write_text('{"version": "0.2.13", "last_updated": "2024-01-01"}')
+
+        return legacy_dir
+
+    def verify_xdg_migration_complete(self, base_dir: Path, legacy_dir: Path):
+        """Verify that migration to XDG structure is complete and correct."""
+        config_dir = base_dir / ".config" / "llama-stack"
+        data_dir = base_dir / ".local" / "share" / "llama-stack"
+        state_dir = base_dir / ".local" / "state" / "llama-stack"
+
+        # Verify distributions moved to config
+        self.assertTrue((config_dir / "distributions").exists())
+        self.assertTrue((config_dir / "distributions" / "ollama").exists())
+        self.assertTrue((config_dir / "distributions" / "ollama" / "ollama-run.yaml").exists())
+
+        # Verify YAML content is preserved
+        yaml_content = (config_dir / "distributions" / "ollama" / "ollama-run.yaml").read_text()
+        self.assertIn("version: 2", yaml_content)
+        self.assertIn("remote::ollama", yaml_content)
+
+        # Verify providers.d moved to config
+        self.assertTrue((config_dir / "providers.d").exists())
+        self.assertTrue((config_dir / "providers.d" / "remote" / "inference").exists())
+        self.assertTrue((config_dir / "providers.d" / "remote" / "inference" / "custom-inference.yaml").exists())
+
+        # Verify checkpoints moved to data
+        self.assertTrue((data_dir / "checkpoints").exists())
+        self.assertTrue((data_dir / "checkpoints" / "meta-llama" / "Llama-3.2-1B-Instruct").exists())
+        self.assertTrue(
+            (data_dir / "checkpoints" / "meta-llama" / "Llama-3.2-1B-Instruct" / "consolidated.00.pth").exists()
+        )
+
+        # Verify model file content preserved
+        model_file = data_dir / "checkpoints" / "meta-llama" / "Llama-3.2-1B-Instruct" / "consolidated.00.pth"
+        self.assertGreater(model_file.stat().st_size, 1000)  # Should be substantial size
+
+        # Verify runtime moved to state
+        self.assertTrue((state_dir / "runtime").exists())
+        self.assertTrue((state_dir / "runtime" / "trace_store.db").exists())
+        self.assertTrue((state_dir / "runtime" / "agent_sessions.db").exists())
+
+        # Verify database files preserved
+        db_file = state_dir / "runtime" / "trace_store.db"
+        db_content = db_file.read_text()
+        self.assertIn("SQLite format 3", db_content)
+
+    def test_fresh_installation_xdg_compliance(self):
+        """Test fresh installation uses XDG-compliant paths."""
+        self.clear_env_vars()
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            # Set custom XDG paths
+            os.environ["XDG_CONFIG_HOME"] = str(base_dir / "custom_config")
+            os.environ["XDG_DATA_HOME"] = str(base_dir / "custom_data")
+            os.environ["XDG_STATE_HOME"] = str(base_dir / "custom_state")
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+                # Mock that no legacy directory exists
+                with patch("llama_stack.distribution.utils.xdg_utils.Path.exists") as mock_exists:
+                    mock_exists.return_value = False
+
+                    # Fresh installation should use XDG paths
+                    config_dir = get_llama_stack_config_dir()
+                    data_dir = get_llama_stack_data_dir()
+                    state_dir = get_llama_stack_state_dir()
+
+                    self.assertEqual(config_dir, base_dir / "custom_config" / "llama-stack")
+                    self.assertEqual(data_dir, base_dir / "custom_data" / "llama-stack")
+                    self.assertEqual(state_dir, base_dir / "custom_state" / "llama-stack")
+
+    def test_complete_migration_workflow(self):
+        """Test complete migration workflow from legacy to XDG."""
+        self.clear_env_vars()
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+
+                # Create realistic legacy structure
+                legacy_dir = self.create_realistic_legacy_structure(base_dir)
+
+                # Verify legacy structure exists
+                self.assertTrue(legacy_dir.exists())
+                self.assertTrue((legacy_dir / "distributions" / "ollama" / "ollama-run.yaml").exists())
+
+                # Perform migration
+                from llama_stack.cli.migrate_xdg import migrate_to_xdg
+
+                with patch("builtins.input") as mock_input:
+                    mock_input.side_effect = ["y", "y"]  # Confirm migration and cleanup
+
+                    result = migrate_to_xdg(dry_run=False)
+                    self.assertTrue(result)
+
+                # Verify migration completed successfully
+                self.verify_xdg_migration_complete(base_dir, legacy_dir)
+
+                # Verify legacy directory was removed
+                self.assertFalse(legacy_dir.exists())
+
+    def test_migration_preserves_file_integrity(self):
+        """Test that migration preserves file integrity and permissions."""
+        self.clear_env_vars()
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+
+                # Create legacy structure
+                legacy_dir = self.create_realistic_legacy_structure(base_dir)
+
+                # Set specific permissions on files
+                config_file = legacy_dir / "distributions" / "ollama" / "ollama-run.yaml"
+                config_file.chmod(0o600)
+                original_stat = config_file.stat()
+
+                # Perform migration
+                from llama_stack.cli.migrate_xdg import migrate_to_xdg
+
+                with patch("builtins.input") as mock_input:
+                    mock_input.side_effect = ["y", "y"]
+                    result = migrate_to_xdg(dry_run=False)
+                    self.assertTrue(result)
+
+                # Verify file integrity
+                migrated_config = base_dir / ".config" / "llama-stack" / "distributions" / "ollama" / "ollama-run.yaml"
+                self.assertTrue(migrated_config.exists())
+
+                # Verify content is identical
+                migrated_content = migrated_config.read_text()
+                self.assertIn("version: 2", migrated_content)
+                self.assertIn("remote::ollama", migrated_content)
+
+                # Verify permissions preserved
+                migrated_stat = migrated_config.stat()
+                self.assertEqual(original_stat.st_mode, migrated_stat.st_mode)
+
+    def test_mixed_legacy_and_xdg_environment(self):
+        """Test behavior in mixed legacy and XDG environment."""
+        self.clear_env_vars()
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            # Set partial XDG environment
+            os.environ["XDG_CONFIG_HOME"] = str(base_dir / "xdg_config")
+            # Leave XDG_DATA_HOME and XDG_STATE_HOME unset
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+
+                # Create legacy directory
+                legacy_dir = self.create_realistic_legacy_structure(base_dir)
+
+                # Should use legacy directory since it exists
+                config_dir = get_llama_stack_config_dir()
+                data_dir = get_llama_stack_data_dir()
+                state_dir = get_llama_stack_state_dir()
+
+                self.assertEqual(config_dir, legacy_dir)
+                self.assertEqual(data_dir, legacy_dir)
+                self.assertEqual(state_dir, legacy_dir)
+
+    def test_template_rendering_with_xdg_paths(self):
+        """Test that templates render correctly with XDG paths."""
+        self.clear_env_vars()
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            # Set XDG environment
+            os.environ["XDG_STATE_HOME"] = str(base_dir / "state")
+            os.environ["XDG_DATA_HOME"] = str(base_dir / "data")
+
+            # Mock shell environment variable expansion
+            def mock_env_expand(template_string):
+                """Mock shell environment variable expansion."""
+                result = template_string
+                result = result.replace("${env.XDG_STATE_HOME:-~/.local/state}", str(base_dir / "state"))
+                result = result.replace("${env.XDG_DATA_HOME:-~/.local/share}", str(base_dir / "data"))
+                return result
+
+            # Test template patterns
+            template_patterns = [
+                "${env.XDG_STATE_HOME:-~/.local/state}/llama-stack/distributions/ollama",
+                "${env.XDG_DATA_HOME:-~/.local/share}/llama-stack/distributions/ollama/files",
+            ]
+
+            expected_results = [
+                str(base_dir / "state" / "llama-stack" / "distributions" / "ollama"),
+                str(base_dir / "data" / "llama-stack" / "distributions" / "ollama" / "files"),
+            ]
+
+            for pattern, expected in zip(template_patterns, expected_results, strict=False):
+                with self.subTest(pattern=pattern):
+                    expanded = mock_env_expand(pattern)
+                    self.assertEqual(expanded, expected)
+
+    def test_cli_integration_with_xdg_paths(self):
+        """Test CLI integration works correctly with XDG paths."""
+        self.clear_env_vars()
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+
+                # Create legacy structure
+                legacy_dir = self.create_realistic_legacy_structure(base_dir)
+
+                # Test CLI migrate command
+                import argparse
+
+                from llama_stack.cli.migrate_xdg import MigrateXDG
+
+                parser = argparse.ArgumentParser()
+                subparsers = parser.add_subparsers()
+                migrate_cmd = MigrateXDG.create(subparsers)
+
+                # Test dry-run
+                args = parser.parse_args(["migrate-xdg", "--dry-run"])
+
+                with patch("builtins.print") as mock_print:
+                    result = migrate_cmd._run_migrate_xdg_cmd(args)
+                    self.assertEqual(result, 0)
+
+                    # Should print dry-run information
+                    print_calls = [call[0][0] for call in mock_print.call_args_list]
+                    self.assertTrue(any("Dry run mode" in call for call in print_calls))
+
+                # Legacy directory should still exist after dry-run
+                self.assertTrue(legacy_dir.exists())
+
+    def test_config_dirs_integration_after_migration(self):
+        """Test that config_dirs works correctly after migration."""
+        self.clear_env_vars()
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+
+                # Create and migrate legacy structure
+                self.create_realistic_legacy_structure(base_dir)
+
+                from llama_stack.cli.migrate_xdg import migrate_to_xdg
+
+                with patch("builtins.input") as mock_input:
+                    mock_input.side_effect = ["y", "y"]
+                    result = migrate_to_xdg(dry_run=False)
+                    self.assertTrue(result)
+
+                # Clear module cache to ensure fresh import
+                import sys
+
+                if "llama_stack.distribution.utils.config_dirs" in sys.modules:
+                    del sys.modules["llama_stack.distribution.utils.config_dirs"]
+
+                # Import config_dirs after migration
+                from llama_stack.distribution.utils.config_dirs import (
+                    DEFAULT_CHECKPOINT_DIR,
+                    DISTRIBS_BASE_DIR,
+                    LLAMA_STACK_CONFIG_DIR,
+                    RUNTIME_BASE_DIR,
+                )
+
+                # Should use XDG paths
+                self.assertEqual(LLAMA_STACK_CONFIG_DIR, base_dir / ".config" / "llama-stack")
+                self.assertEqual(DEFAULT_CHECKPOINT_DIR, base_dir / ".local" / "share" / "llama-stack" / "checkpoints")
+                self.assertEqual(RUNTIME_BASE_DIR, base_dir / ".local" / "state" / "llama-stack" / "runtime")
+                self.assertEqual(DISTRIBS_BASE_DIR, base_dir / ".config" / "llama-stack" / "distributions")
+
+    def test_real_file_operations_with_xdg_paths(self):
+        """Test real file operations work correctly with XDG paths."""
+        self.clear_env_vars()
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            # Set XDG environment
+            os.environ["XDG_CONFIG_HOME"] = str(base_dir / "config")
+            os.environ["XDG_DATA_HOME"] = str(base_dir / "data")
+            os.environ["XDG_STATE_HOME"] = str(base_dir / "state")
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+                with patch("llama_stack.distribution.utils.xdg_utils.Path.exists") as mock_exists:
+                    mock_exists.return_value = False
+
+                    # Get XDG paths
+                    config_dir = get_llama_stack_config_dir()
+                    data_dir = get_llama_stack_data_dir()
+                    state_dir = get_llama_stack_state_dir()
+
+                    # Create directories
+                    config_dir.mkdir(parents=True)
+                    data_dir.mkdir(parents=True)
+                    state_dir.mkdir(parents=True)
+
+                    # Test writing configuration files
+                    config_file = config_dir / "test_config.yaml"
+                    config_data = {"version": "2", "test": True}
+
+                    with open(config_file, "w") as f:
+                        yaml.dump(config_data, f)
+
+                    # Test reading configuration files
+                    with open(config_file) as f:
+                        loaded_config = yaml.safe_load(f)
+
+                    self.assertEqual(loaded_config, config_data)
+
+                    # Test creating nested directory structure
+                    model_dir = data_dir / "checkpoints" / "meta-llama" / "test-model"
+                    model_dir.mkdir(parents=True)
+
+                    # Test writing large files
+                    model_file = model_dir / "model.bin"
+                    test_data = b"test model data" * 1000
+                    model_file.write_bytes(test_data)
+
+                    # Verify file integrity
+                    read_data = model_file.read_bytes()
+                    self.assertEqual(read_data, test_data)
+
+                    # Test state files
+                    state_file = state_dir / "runtime" / "session.db"
+                    state_file.parent.mkdir(parents=True)
+                    state_file.write_text("SQLite format 3\x00test database")
+
+                    # Verify state file
+                    state_content = state_file.read_text()
+                    self.assertIn("SQLite format 3", state_content)
+
+    def test_backwards_compatibility_scenario(self):
+        """Test complete backwards compatibility scenario."""
+        self.clear_env_vars()
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            # Scenario: User has existing legacy installation
+            legacy_dir = self.create_realistic_legacy_structure(base_dir)
+
+            # User sets legacy environment variable
+            os.environ["LLAMA_STACK_CONFIG_DIR"] = str(legacy_dir)
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+
+                # Should continue using legacy paths
+                config_dir = get_llama_stack_config_dir()
+                data_dir = get_llama_stack_data_dir()
+                state_dir = get_llama_stack_state_dir()
+
+                self.assertEqual(config_dir, legacy_dir)
+                self.assertEqual(data_dir, legacy_dir)
+                self.assertEqual(state_dir, legacy_dir)
+
+                # Should be able to access existing files
+                yaml_file = legacy_dir / "distributions" / "ollama" / "ollama-run.yaml"
+                self.assertTrue(yaml_file.exists())
+
+                # Should be able to parse existing configuration
+                with open(yaml_file) as f:
+                    config = yaml.safe_load(f)
+
+                self.assertEqual(config["version"], 2)
+
+    def test_error_recovery_scenarios(self):
+        """Test error recovery in various scenarios."""
+        self.clear_env_vars()
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+
+                # Scenario 1: Partial migration failure
+                self.create_realistic_legacy_structure(base_dir)
+
+                # Create conflicting file in target location
+                config_dir = base_dir / ".config" / "llama-stack"
+                config_dir.mkdir(parents=True)
+
+                conflicting_file = config_dir / "distributions"
+                conflicting_file.touch()  # Create file instead of directory
+
+                from llama_stack.cli.migrate_xdg import migrate_to_xdg
+
+                with patch("builtins.input") as mock_input:
+                    mock_input.side_effect = ["y", "n"]  # Confirm migration, don't cleanup
+
+                    with patch("builtins.print") as mock_print:
+                        result = migrate_to_xdg(dry_run=False)
+
+                        # Should handle conflicts gracefully
+                        print_calls = [call[0][0] for call in mock_print.call_args_list]
+                        conflict_mentioned = any(
+                            "Warning" in call or "conflict" in call.lower() for call in print_calls
+                        )
+
+                        # Migration should complete with warnings
+                        self.assertTrue(result or conflict_mentioned)
+
+    def test_cross_platform_compatibility(self):
+        """Test cross-platform compatibility of XDG implementation."""
+        self.clear_env_vars()
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            # Test with different path separators and formats
+            if os.name == "nt":  # Windows
+                # Test Windows-style paths
+                os.environ["XDG_CONFIG_HOME"] = str(base_dir / "config").replace("/", "\\")
+                os.environ["XDG_DATA_HOME"] = str(base_dir / "data").replace("/", "\\")
+            else:  # Unix-like
+                # Test Unix-style paths
+                os.environ["XDG_CONFIG_HOME"] = str(base_dir / "config")
+                os.environ["XDG_DATA_HOME"] = str(base_dir / "data")
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+                with patch("llama_stack.distribution.utils.xdg_utils.Path.exists") as mock_exists:
+                    mock_exists.return_value = False
+
+                    # Should work regardless of platform
+                    config_dir = get_llama_stack_config_dir()
+                    data_dir = get_llama_stack_data_dir()
+
+                    # Paths should be valid for the current platform
+                    self.assertTrue(config_dir.is_absolute())
+                    self.assertTrue(data_dir.is_absolute())
+
+                    # Should be able to create directories
+                    config_dir.mkdir(parents=True, exist_ok=True)
+                    data_dir.mkdir(parents=True, exist_ok=True)
+
+                    self.assertTrue(config_dir.exists())
+                    self.assertTrue(data_dir.exists())
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/integration/test_xdg_migration.py

@@ -0,0 +1,516 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import os
+import tempfile
+import unittest
+from pathlib import Path
+from unittest.mock import patch
+
+from llama_stack.cli.migrate_xdg import migrate_to_xdg
+
+
+class TestXDGMigrationIntegration(unittest.TestCase):
+    """Integration tests for XDG migration functionality."""
+
+    def setUp(self):
+        """Set up test environment."""
+        # Store original environment variables
+        self.original_env = {}
+        for key in ["XDG_CONFIG_HOME", "XDG_DATA_HOME", "XDG_STATE_HOME", "LLAMA_STACK_CONFIG_DIR"]:
+            self.original_env[key] = os.environ.get(key)
+
+        # Clear environment variables
+        for key in self.original_env:
+            os.environ.pop(key, None)
+
+    def tearDown(self):
+        """Clean up test environment."""
+        # Restore original environment variables
+        for key, value in self.original_env.items():
+            if value is None:
+                os.environ.pop(key, None)
+            else:
+                os.environ[key] = value
+
+    def create_legacy_structure(self, base_dir: Path) -> Path:
+        """Create a realistic legacy ~/.llama directory structure."""
+        legacy_dir = base_dir / ".llama"
+        legacy_dir.mkdir()
+
+        # Create distributions
+        distributions_dir = legacy_dir / "distributions"
+        distributions_dir.mkdir()
+
+        # Create sample distribution
+        ollama_dir = distributions_dir / "ollama"
+        ollama_dir.mkdir()
+        (ollama_dir / "ollama-run.yaml").write_text("version: 2\napis: []\n")
+        (ollama_dir / "build.yaml").write_text("name: ollama\n")
+
+        # Create providers.d
+        providers_dir = legacy_dir / "providers.d"
+        providers_dir.mkdir()
+        (providers_dir / "remote").mkdir()
+        (providers_dir / "remote" / "inference").mkdir()
+        (providers_dir / "remote" / "inference" / "custom.yaml").write_text("provider_type: remote::custom\n")
+
+        # Create checkpoints
+        checkpoints_dir = legacy_dir / "checkpoints"
+        checkpoints_dir.mkdir()
+        model_dir = checkpoints_dir / "meta-llama" / "Llama-3.2-1B-Instruct"
+        model_dir.mkdir(parents=True)
+        (model_dir / "consolidated.00.pth").write_text("fake model weights")
+        (model_dir / "params.json").write_text('{"dim": 2048}')
+
+        # Create runtime
+        runtime_dir = legacy_dir / "runtime"
+        runtime_dir.mkdir()
+        (runtime_dir / "trace_store.db").write_text("fake sqlite database")
+
+        # Create some fake files in various subdirectories
+        (legacy_dir / "config.json").write_text('{"version": "0.2.13"}')
+
+        return legacy_dir
+
+    def verify_xdg_structure(self, base_dir: Path, legacy_dir: Path):
+        """Verify that the XDG structure was created correctly."""
+        config_dir = base_dir / ".config" / "llama-stack"
+        data_dir = base_dir / ".local" / "share" / "llama-stack"
+        state_dir = base_dir / ".local" / "state" / "llama-stack"
+
+        # Verify distributions moved to config
+        self.assertTrue((config_dir / "distributions").exists())
+        self.assertTrue((config_dir / "distributions" / "ollama" / "ollama-run.yaml").exists())
+
+        # Verify providers.d moved to config
+        self.assertTrue((config_dir / "providers.d").exists())
+        self.assertTrue((config_dir / "providers.d" / "remote" / "inference" / "custom.yaml").exists())
+
+        # Verify checkpoints moved to data
+        self.assertTrue((data_dir / "checkpoints").exists())
+        self.assertTrue(
+            (data_dir / "checkpoints" / "meta-llama" / "Llama-3.2-1B-Instruct" / "consolidated.00.pth").exists()
+        )
+
+        # Verify runtime moved to state
+        self.assertTrue((state_dir / "runtime").exists())
+        self.assertTrue((state_dir / "runtime" / "trace_store.db").exists())
+
+    def test_full_migration_workflow(self):
+        """Test complete migration workflow from legacy to XDG."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            # Set up fake home directory
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+
+                # Create legacy structure
+                legacy_dir = self.create_legacy_structure(base_dir)
+
+                # Verify legacy structure exists
+                self.assertTrue(legacy_dir.exists())
+                self.assertTrue((legacy_dir / "distributions").exists())
+                self.assertTrue((legacy_dir / "checkpoints").exists())
+
+                # Perform migration with user confirmation
+                with patch("builtins.input") as mock_input:
+                    mock_input.side_effect = ["y", "y"]  # Confirm migration and cleanup
+
+                    result = migrate_to_xdg(dry_run=False)
+                    self.assertTrue(result)
+
+                # Verify XDG structure was created
+                self.verify_xdg_structure(base_dir, legacy_dir)
+
+                # Verify legacy directory was removed
+                self.assertFalse(legacy_dir.exists())
+
+    def test_migration_dry_run(self):
+        """Test dry run migration (no actual file movement)."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+
+                # Create legacy structure
+                legacy_dir = self.create_legacy_structure(base_dir)
+
+                # Perform dry run migration
+                with patch("builtins.print") as mock_print:
+                    result = migrate_to_xdg(dry_run=True)
+                    self.assertTrue(result)
+
+                    # Check that dry run message was printed
+                    print_calls = [call[0][0] for call in mock_print.call_args_list]
+                    self.assertTrue(any("Dry run mode" in call for call in print_calls))
+
+                # Verify nothing was actually moved
+                self.assertTrue(legacy_dir.exists())
+                self.assertTrue((legacy_dir / "distributions").exists())
+                self.assertFalse((base_dir / ".config" / "llama-stack").exists())
+
+    def test_migration_user_cancellation(self):
+        """Test migration when user cancels."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+
+                # Create legacy structure
+                legacy_dir = self.create_legacy_structure(base_dir)
+
+                # User cancels migration
+                with patch("builtins.input") as mock_input:
+                    mock_input.return_value = "n"
+
+                    result = migrate_to_xdg(dry_run=False)
+                    self.assertFalse(result)
+
+                # Verify nothing was moved
+                self.assertTrue(legacy_dir.exists())
+                self.assertTrue((legacy_dir / "distributions").exists())
+                self.assertFalse((base_dir / ".config" / "llama-stack").exists())
+
+    def test_migration_with_existing_xdg_directories(self):
+        """Test migration when XDG directories already exist."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+
+                # Create legacy structure
+                legacy_dir = self.create_legacy_structure(base_dir)
+
+                # Create existing XDG structure with conflicting files
+                config_dir = base_dir / ".config" / "llama-stack"
+                config_dir.mkdir(parents=True)
+                (config_dir / "distributions").mkdir()
+                (config_dir / "distributions" / "existing.yaml").write_text("existing config")
+
+                # Perform migration
+                with patch("builtins.input") as mock_input:
+                    mock_input.side_effect = ["y", "n"]  # Confirm migration, don't cleanup
+                    with patch("builtins.print") as mock_print:
+                        result = migrate_to_xdg(dry_run=False)
+                        self.assertTrue(result)
+
+                        # Check that warning was printed
+                        print_calls = [call[0][0] for call in mock_print.call_args_list]
+                        self.assertTrue(any("Warning: Target already exists" in call for call in print_calls))
+
+                # Verify existing file wasn't overwritten
+                self.assertTrue((config_dir / "distributions" / "existing.yaml").exists())
+
+                # Legacy distributions should still exist due to conflict
+                self.assertTrue((legacy_dir / "distributions").exists())
+
+    def test_migration_partial_success(self):
+        """Test migration when some items succeed and others fail."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+
+                # Create legacy structure
+                self.create_legacy_structure(base_dir)
+
+                # Create readonly target directory to simulate permission error
+                config_dir = base_dir / ".config" / "llama-stack"
+                config_dir.mkdir(parents=True)
+                distributions_target = config_dir / "distributions"
+                distributions_target.mkdir()
+                distributions_target.chmod(0o444)  # Read-only
+
+                try:
+                    # Perform migration
+                    with patch("builtins.input") as mock_input:
+                        mock_input.side_effect = ["y", "n"]  # Confirm migration, don't cleanup
+
+                        migrate_to_xdg(dry_run=False)
+                        # Should return True even with partial success
+
+                        # Some items should have been migrated successfully
+                        self.assertTrue((base_dir / ".local" / "share" / "llama-stack" / "checkpoints").exists())
+
+                finally:
+                    # Restore permissions for cleanup
+                    distributions_target.chmod(0o755)
+
+    def test_migration_empty_legacy_directory(self):
+        """Test migration when legacy directory exists but is empty."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+
+                # Create empty legacy directory
+                legacy_dir = base_dir / ".llama"
+                legacy_dir.mkdir()
+
+                result = migrate_to_xdg(dry_run=False)
+                self.assertTrue(result)
+
+    def test_migration_preserves_file_permissions(self):
+        """Test that migration preserves file permissions."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+
+                # Create legacy structure with specific permissions
+                legacy_dir = self.create_legacy_structure(base_dir)
+
+                # Set specific permissions on a file
+                config_file = legacy_dir / "distributions" / "ollama" / "ollama-run.yaml"
+                config_file.chmod(0o600)
+                original_stat = config_file.stat()
+
+                # Perform migration
+                with patch("builtins.input") as mock_input:
+                    mock_input.side_effect = ["y", "y"]
+                    result = migrate_to_xdg(dry_run=False)
+                    self.assertTrue(result)
+
+                # Verify permissions were preserved
+                migrated_file = base_dir / ".config" / "llama-stack" / "distributions" / "ollama" / "ollama-run.yaml"
+                self.assertTrue(migrated_file.exists())
+                migrated_stat = migrated_file.stat()
+                self.assertEqual(original_stat.st_mode, migrated_stat.st_mode)
+
+    def test_migration_preserves_directory_structure(self):
+        """Test that migration preserves complex directory structures."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+
+                # Create complex legacy structure
+                legacy_dir = base_dir / ".llama"
+                legacy_dir.mkdir()
+
+                # Create nested structure
+                complex_path = legacy_dir / "checkpoints" / "org" / "model" / "variant" / "files"
+                complex_path.mkdir(parents=True)
+                (complex_path / "model.bin").write_text("model data")
+                (complex_path / "config.json").write_text("config data")
+
+                # Perform migration
+                with patch("builtins.input") as mock_input:
+                    mock_input.side_effect = ["y", "y"]
+                    result = migrate_to_xdg(dry_run=False)
+                    self.assertTrue(result)
+
+                # Verify structure was preserved
+                migrated_path = (
+                    base_dir
+                    / ".local"
+                    / "share"
+                    / "llama-stack"
+                    / "checkpoints"
+                    / "org"
+                    / "model"
+                    / "variant"
+                    / "files"
+                )
+                self.assertTrue(migrated_path.exists())
+                self.assertTrue((migrated_path / "model.bin").exists())
+                self.assertTrue((migrated_path / "config.json").exists())
+
+    def test_migration_with_symlinks(self):
+        """Test migration with symbolic links in legacy directory."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+
+                # Create legacy structure
+                legacy_dir = self.create_legacy_structure(base_dir)
+
+                # Create a symlink
+                actual_file = legacy_dir / "actual_config.yaml"
+                actual_file.write_text("actual config content")
+
+                symlink_file = legacy_dir / "distributions" / "symlinked_config.yaml"
+                symlink_file.symlink_to(actual_file)
+
+                # Perform migration
+                with patch("builtins.input") as mock_input:
+                    mock_input.side_effect = ["y", "y"]
+                    result = migrate_to_xdg(dry_run=False)
+                    self.assertTrue(result)
+
+                # Verify symlink was preserved
+                migrated_symlink = base_dir / ".config" / "llama-stack" / "distributions" / "symlinked_config.yaml"
+                self.assertTrue(migrated_symlink.exists())
+                self.assertTrue(migrated_symlink.is_symlink())
+
+    def test_migration_large_files(self):
+        """Test migration with large files (simulated)."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+
+                # Create legacy structure
+                legacy_dir = self.create_legacy_structure(base_dir)
+
+                # Create a larger file (1MB)
+                large_file = legacy_dir / "checkpoints" / "large_model.bin"
+                large_file.write_bytes(b"0" * (1024 * 1024))
+
+                # Perform migration
+                with patch("builtins.input") as mock_input:
+                    mock_input.side_effect = ["y", "y"]
+                    result = migrate_to_xdg(dry_run=False)
+                    self.assertTrue(result)
+
+                # Verify large file was moved correctly
+                migrated_file = base_dir / ".local" / "share" / "llama-stack" / "checkpoints" / "large_model.bin"
+                self.assertTrue(migrated_file.exists())
+                self.assertEqual(migrated_file.stat().st_size, 1024 * 1024)
+
+    def test_migration_with_unicode_filenames(self):
+        """Test migration with unicode filenames."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+
+                # Create legacy structure with unicode filenames
+                legacy_dir = base_dir / ".llama"
+                legacy_dir.mkdir()
+
+                unicode_dir = legacy_dir / "distributions" / "配置"
+                unicode_dir.mkdir(parents=True)
+                unicode_file = unicode_dir / "模型.yaml"
+                unicode_file.write_text("unicode content")
+
+                # Perform migration
+                with patch("builtins.input") as mock_input:
+                    mock_input.side_effect = ["y", "y"]
+                    result = migrate_to_xdg(dry_run=False)
+                    self.assertTrue(result)
+
+                # Verify unicode files were migrated
+                migrated_file = base_dir / ".config" / "llama-stack" / "distributions" / "配置" / "模型.yaml"
+                self.assertTrue(migrated_file.exists())
+                self.assertEqual(migrated_file.read_text(), "unicode content")
+
+
+class TestXDGMigrationCLI(unittest.TestCase):
+    """Test the CLI interface for XDG migration."""
+
+    def setUp(self):
+        """Set up test environment."""
+        self.original_env = {}
+        for key in ["XDG_CONFIG_HOME", "XDG_DATA_HOME", "XDG_STATE_HOME", "LLAMA_STACK_CONFIG_DIR"]:
+            self.original_env[key] = os.environ.get(key)
+            os.environ.pop(key, None)
+
+    def tearDown(self):
+        """Clean up test environment."""
+        for key, value in self.original_env.items():
+            if value is None:
+                os.environ.pop(key, None)
+            else:
+                os.environ[key] = value
+
+    def test_cli_migrate_command_exists(self):
+        """Test that the migrate-xdg CLI command is properly registered."""
+        from llama_stack.cli.llama import LlamaCLIParser
+
+        parser = LlamaCLIParser()
+
+        # Parse help to see if migrate-xdg is listed
+        with patch("sys.argv", ["llama", "--help"]):
+            with patch("sys.exit"):
+                with patch("builtins.print") as mock_print:
+                    try:
+                        parser.parse_args()
+                    except SystemExit:
+                        pass
+
+                    # Check if migrate-xdg appears in help output
+                    help_output = "\n".join([call[0][0] for call in mock_print.call_args_list])
+                    self.assertIn("migrate-xdg", help_output)
+
+    def test_cli_migrate_dry_run(self):
+        """Test CLI migrate command with dry-run flag."""
+        import argparse
+
+        from llama_stack.cli.migrate_xdg import MigrateXDG
+
+        # Create parser and add migrate command
+        parser = argparse.ArgumentParser()
+        subparsers = parser.add_subparsers()
+        MigrateXDG.create(subparsers)
+
+        # Test dry-run flag
+        args = parser.parse_args(["migrate-xdg", "--dry-run"])
+        self.assertTrue(args.dry_run)
+
+        # Test without dry-run flag
+        args = parser.parse_args(["migrate-xdg"])
+        self.assertFalse(args.dry_run)
+
+    def test_cli_migrate_execution(self):
+        """Test CLI migrate command execution."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            base_dir = Path(temp_dir)
+
+            with patch("llama_stack.distribution.utils.xdg_utils.Path.home") as mock_home:
+                mock_home.return_value = base_dir
+
+                # Create legacy directory
+                legacy_dir = base_dir / ".llama"
+                legacy_dir.mkdir()
+                (legacy_dir / "test_file").touch()
+
+                import argparse
+
+                from llama_stack.cli.migrate_xdg import MigrateXDG
+
+                # Create parser and command
+                parser = argparse.ArgumentParser()
+                subparsers = parser.add_subparsers()
+                migrate_cmd = MigrateXDG.create(subparsers)
+
+                # Parse arguments
+                args = parser.parse_args(["migrate-xdg", "--dry-run"])
+
+                # Execute command
+                with patch("builtins.print") as mock_print:
+                    migrate_cmd._run_migrate_xdg_cmd(args)
+
+                    # Verify output was printed
+                    print_calls = [call[0][0] for call in mock_print.call_args_list]
+                    self.assertTrue(any("Found legacy directory" in call for call in print_calls))
+
+
+if __name__ == "__main__":
+    unittest.main()