From b57cee65cc7b54136a3027a036445b601e5c2756 Mon Sep 17 00:00:00 2001 From: yassin Date: Sat, 5 Sep 2026 00:22:42 +0000 Subject: [PATCH] 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 | 19 ++++++++++-- .../test_litellm/proxy/client/cli/test_pi.py | 30 +++++++++++++++++++ 2 files changed, 46 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/client/cli/commands/pi.py b/litellm/proxy/client/cli/commands/pi.py index 668b803a33d..b3b9d520a5a 100644 --- a/litellm/proxy/client/cli/commands/pi.py +++ b/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 diff --git a/tests/test_litellm/proxy/client/cli/test_pi.py b/tests/test_litellm/proxy/client/cli/test_pi.py index ee3222a77e3..68c0ac70064 100644 --- a/tests/test_litellm/proxy/client/cli/test_pi.py +++ b/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")