Merge pull request #1857 from transformerlab/fix/cli-config-atomic-write

fix(cli): prevent ~/.lab/config.json from being silently wiped

Author: aliasaria

cli/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "transformerlab-cli"
-version = "0.0.21"
+version = "0.0.22"
 description = "Transformer Lab CLI"
 requires-python = ">=3.10"
 authors = [{ name = "Transformer Lab", email = "hello@transformerlab.ai" }]

cli/src/transformerlab_cli/util/config.py

@@ -1,5 +1,6 @@
 import os
 import json
+import time
 from typing import Any
 from urllib.parse import urlparse
 
@@ -21,7 +22,12 @@
 
 
 def load_config() -> dict[str, Any]:
-    """Load config from file, return empty dict if not found."""
+    """Load config from file, return empty dict if not found.
+
+    If the file exists but cannot be parsed (corrupt/truncated), move it
+    aside to config.json.corrupt-<ts> and return {} so the next write
+    does not silently overwrite a recoverable file.
+    """
     global cached_config
 
     if cached_config is not None:
@@ -33,21 +39,47 @@ def load_config() -> dict[str, Any]:
         with open(CONFIG_FILE, "r", encoding="utf-8") as f:
             cached_config = json.loads(f.read())
         return cached_config
-    except (json.JSONDecodeError, OSError):
+    except (json.JSONDecodeError, OSError) as e:
+        backup_path = f"{CONFIG_FILE}.corrupt-{int(time.time())}"
+        try:
+            os.rename(CONFIG_FILE, backup_path)
+            console.print(
+                f"[error]Error:[/error] Config file is unreadable ({e}). "
+                f"Moved to [label]{backup_path}[/label] to avoid overwriting it."
+            )
+        except OSError as rename_err:
+            console.print(
+                f"[error]Error:[/error] Config file is unreadable ({e}) and could not be "
+                f"backed up ({rename_err}). Refusing to continue."
+            )
+            raise typer.Exit(1)
         return {}
 
 
 def _save_config(config: dict[str, Any]) -> bool:
-    """Save config to file. Returns True on success."""
+    """Save config to file atomically. Returns True on success.
+
+    Writes to a sibling tmp file and renames over the target so readers
+    never observe a truncated or partially-written config.
+    """
     global cached_config
+    tmp_path = f"{CONFIG_FILE}.tmp"
     try:
         os.makedirs(CONFIG_DIR, exist_ok=True)
-        with open(CONFIG_FILE, "w", encoding="utf-8") as f:
+        with open(tmp_path, "w", encoding="utf-8") as f:
             f.write(json.dumps(config, indent=2))
+            f.flush()
+            os.fsync(f.fileno())
+        os.replace(tmp_path, CONFIG_FILE)
         cached_config = config
         return True
     except OSError as e:
         console.print(f"[error]Error:[/error] Failed to save config: {e}")
+        try:
+            if os.path.exists(tmp_path):
+                os.unlink(tmp_path)
+        except OSError:
+            pass
         return False
 
 

cli/tests/commands/test_config.py

@@ -1,6 +1,7 @@
 """Tests for config command and utility functions."""
 
 import json
+import os
 from typer.testing import CliRunner
 from transformerlab_cli.main import app
 from transformerlab_cli.util.config import (
@@ -137,6 +138,54 @@ def test_load_config_with_data(tmp_config_dir):
         assert config["server"] == "http://test:8338"
 
 
+def test_load_config_corrupt_moves_aside_and_returns_empty(tmp_config_dir):
+    """A corrupt config must not silently resolve to {}; it must be renamed aside."""
+    import glob
+
+    config_dir, config_file = tmp_config_dir
+    with open(config_file, "w", encoding="utf-8") as f:
+        f.write("{not: valid json,,")
+    with _patch_config_paths(config_dir, config_file):
+        config = load_config()
+        assert config == {}
+    assert not os.path.exists(config_file), "corrupt config should have been renamed aside"
+    backups = glob.glob(f"{config_file}.corrupt-*")
+    assert len(backups) == 1, f"expected one backup, got {backups}"
+
+
+def test_set_config_does_not_wipe_existing_keys_when_file_is_fine(tmp_config_dir):
+    """Regression: a set_config call must merge, never wipe unrelated keys."""
+    config_dir, config_file = tmp_config_dir
+    with open(config_file, "w", encoding="utf-8") as f:
+        f.write(
+            json.dumps(
+                {
+                    "server": "http://localhost:8338",
+                    "team_id": "abc-123",
+                    "user_email": "u@example.com",
+                    "current_experiment": "alpha",
+                }
+            )
+        )
+    with _patch_config_paths(config_dir, config_file):
+        assert set_config("current_experiment", "beta") is True
+        with open(config_file, "r", encoding="utf-8") as f:
+            data = json.loads(f.read())
+    assert data["server"] == "http://localhost:8338"
+    assert data["team_id"] == "abc-123"
+    assert data["user_email"] == "u@example.com"
+    assert data["current_experiment"] == "beta"
+
+
+def test_save_config_is_atomic_no_tmp_left_behind(tmp_config_dir):
+    """After a successful save, the sibling .tmp file must not exist."""
+    config_dir, config_file = tmp_config_dir
+    with _patch_config_paths(config_dir, config_file):
+        assert set_config("server", "http://localhost:8338") is True
+    assert os.path.exists(config_file)
+    assert not os.path.exists(f"{config_file}.tmp")
+
+
 # --- JSON format output ---
 
 

cli/tests/commands/test_login.py

@@ -50,3 +50,35 @@ def test_logout_success(_mock_delete):
     """Test successful logout."""
     result = runner.invoke(app, ["logout"])
     assert result.exit_code == 0
+
+
+@patch("transformerlab_cli.commands.logout.delete_api_key", return_value=True)
+def test_logout_clears_user_and_team_keys_but_not_server(_mock_delete):
+    """Regression: logout must only clear session-scoped keys, keep server.
+
+    And, via the autouse isolation fixture, must never touch the real
+    ~/.lab/config.json on the developer's machine.
+    """
+    import json
+
+    from transformerlab_cli.util.shared import CONFIG_FILE
+
+    with open(CONFIG_FILE, "w", encoding="utf-8") as f:
+        f.write(
+            json.dumps(
+                {
+                    "server": "http://localhost:8338",
+                    "team_id": "abc-123",
+                    "team_name": "Team Name",
+                    "user_email": "user@example.com",
+                    "current_experiment": "alpha",
+                }
+            )
+        )
+
+    result = runner.invoke(app, ["logout"])
+    assert result.exit_code == 0
+
+    with open(CONFIG_FILE, "r", encoding="utf-8") as f:
+        data = json.loads(f.read())
+    assert data == {"server": "http://localhost:8338"}, data

cli/tests/conftest.py

@@ -7,6 +7,40 @@
 import pytest
 
 
+@pytest.fixture(autouse=True)
+def _isolate_lab_config_from_user_home(tmp_path, monkeypatch):
+    """Point CONFIG_DIR/CONFIG_FILE and CREDENTIALS_DIR/CREDENTIALS_FILE at a
+    per-test tmp directory so tests cannot read from — or wipe — the real
+    ``~/.lab/config.json`` / ``~/.lab/credentials`` on the developer's machine.
+
+    Applied to every test via ``autouse=True``; any test that needs the raw
+    paths can still accept the ``tmp_config_dir`` fixture explicitly.
+    """
+    fake_lab_dir = tmp_path / "_isolated_lab"
+    fake_lab_dir.mkdir(exist_ok=True)
+    fake_config_file = str(fake_lab_dir / "config.json")
+    fake_credentials_file = str(fake_lab_dir / "credentials")
+
+    import transformerlab_cli.util.config as config_mod
+    import transformerlab_cli.util.shared as shared_mod
+
+    monkeypatch.setattr(shared_mod, "CONFIG_DIR", str(fake_lab_dir))
+    monkeypatch.setattr(shared_mod, "CONFIG_FILE", fake_config_file)
+    monkeypatch.setattr(shared_mod, "CREDENTIALS_DIR", str(fake_lab_dir))
+    monkeypatch.setattr(shared_mod, "CREDENTIALS_FILE", fake_credentials_file)
+    monkeypatch.setattr(config_mod, "CONFIG_DIR", str(fake_lab_dir))
+    monkeypatch.setattr(config_mod, "CONFIG_FILE", fake_config_file)
+    monkeypatch.setattr(config_mod, "cached_config", None)
+
+    try:
+        import transformerlab_cli.util.auth as auth_mod
+
+        monkeypatch.setattr(auth_mod, "CREDENTIALS_DIR", str(fake_lab_dir), raising=False)
+        monkeypatch.setattr(auth_mod, "CREDENTIALS_FILE", fake_credentials_file, raising=False)
+    except ImportError:
+        pass
+
+
 @pytest.fixture()
 def tmp_config_dir(tmp_path):
     """Provide a temporary config directory and patch CONFIG_DIR/CONFIG_FILE."""