fix(cli): write pi models atomically and privately

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/client/cli/commands/pi.py

@@ -6,6 +6,8 @@ the short-lived login token never lands on disk.
 """
 
 import json
+import os
+import tempfile
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass
 from pathlib import Path
@@ -156,13 +158,24 @@ def sync_models_json(
         **current,
         "providers": {**existing_providers, PI_PROVIDER_NAME: provider_block(base_url, model_ids, limits)},
     }
-    staging: Final = path.with_name(path.name + ".tmp")
     try:
         path.parent.mkdir(parents=True, exist_ok=True)
-        staging.write_text(json.dumps(updated, indent=2) + "\n")
-        staging.replace(path)
     except OSError as e:
         return PiSyncError(f"Could not write {path}: {e}")
+    try:
+        fd, tmp_name = tempfile.mkstemp(dir=path.parent, prefix=path.name + ".", suffix=".tmp")
+    except OSError as e:
+        return PiSyncError(f"Could not write {path}: {e}")
+    try:
+        with os.fdopen(fd, "w") as file:
+            file.write(json.dumps(updated, indent=2) + "\n")
+        os.replace(tmp_name, path)
+    except OSError as e:
+        try:
+            os.unlink(tmp_name)
+        except FileNotFoundError:
+            pass
+        return PiSyncError(f"Could not write {path}: {e}")
     return None
 
 

tests/test_litellm/proxy/client/cli/test_pi.py

@@ -1,4 +1,7 @@
 import json
+import os
+import stat
+from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 
 import requests
@@ -179,6 +182,33 @@ class TestSyncModelsJson:
         assert sync_models_json(path, "http://localhost:4000", ("m-1",)) is None
         assert [p.name for p in tmp_path.iterdir()] == ["models.json"]
 
+    def test_written_file_is_private(self, tmp_path):
+        path = tmp_path / "models.json"
+        assert sync_models_json(path, "http://localhost:4000", ("m-1",)) is None
+        if os.name != "nt":
+            assert stat.S_IMODE(path.stat().st_mode) == 0o600
+
+        path.write_text(json.dumps({"providers": {"other": {"apiKey": "literal-secret"}}}))
+        path.chmod(0o644)
+        assert sync_models_json(path, "http://localhost:4000", ("m-2",)) is None
+        if os.name != "nt":
+            assert stat.S_IMODE(path.stat().st_mode) == 0o600
+
+    def test_concurrent_syncs_do_not_collide(self, tmp_path):
+        path = tmp_path / "models.json"
+        model_lists = (("m-a",), ("m-b",))
+
+        def sync(model_ids):
+            return sync_models_json(path, "http://localhost:4000", model_ids)
+
+        with ThreadPoolExecutor(max_workers=2) as executor:
+            results = [result for _ in range(30) for result in executor.map(sync, model_lists)]
+
+        assert results == [None] * 60
+        written = json.loads(path.read_text())
+        assert written["providers"]["litellm"]["models"] in ([{"id": "m-a"}], [{"id": "m-b"}])
+        assert list(tmp_path.glob("models.json.*.tmp")) == []
+
     def test_invalid_json_is_a_value_and_file_untouched(self, tmp_path):
         path = tmp_path / "models.json"
         path.write_text("{not json")