scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-30 01:52:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 70

refactor(proxy): extract _prepare_ui_directory; add symlinks=True; cover read-only branch directly

Browse source 
Extract the restructuring decision into _prepare_ui_directory so tests can
call the actual code path with os.access mocked, rather than replicating the
logic manually. Add symlinks=True to copytree so source symlinks are preserved
instead of being dereferenced into the temp copy. Add two new tests: one that
exercises the is_writable=False branch through _prepare_ui_directory, one that
verifies the temp dir is cleaned up when copytree fails.
This commit is contained in:
Daniele Salvador 2026-06-01 01:05:10 +02:00
parent 1f75cad6fe
commit 82a7bf12a6
2 changed files with 94 additions and 25 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

55
litellm/proxy/proxy_server.py
View file

@ -1659,33 +1659,38 @@ try:
# Another process may have already moved this file.
continue
try:
is_pre_restructured = _is_ui_pre_restructured(ui_path)
is_writable = os.access(ui_path, os.W_OK)
if is_pre_restructured:
def _prepare_ui_directory(path: str) -> str:
"""Return a fully restructured UI path, copying to a temp dir if the source is read-only."""
if _is_ui_pre_restructured(path):
verbose_proxy_logger.info(
f"Skipping UI restructuring: {ui_path} is already pre-restructured"
f"Skipping UI restructuring: {path} is already pre-restructured"
)
elif is_writable:
_restructure_ui_html_files(ui_path)
verbose_proxy_logger.info(f"Restructured UI directory: {ui_path}")
else:
tmp_dir = tempfile.mkdtemp(prefix="litellm_ui_")
try:
shutil.copytree(ui_path, tmp_dir, dirs_exist_ok=True)
_restructure_ui_html_files(tmp_dir)
atexit.register(shutil.rmtree, tmp_dir, True)
ui_path = tmp_dir
verbose_proxy_logger.info(
f"Copied read-only UI to temp dir and restructured: {tmp_dir}"
)
except Exception:
shutil.rmtree(tmp_dir, ignore_errors=True)
verbose_proxy_logger.warning(
f"Failed to copy and restructure UI from {ui_path}; "
f"extensionless routes like /ui/login may return 404"
)
return path
if os.access(path, os.W_OK):
_restructure_ui_html_files(path)
verbose_proxy_logger.info(f"Restructured UI directory: {path}")
return path
tmp_dir = tempfile.mkdtemp(prefix="litellm_ui_")
try:
shutil.copytree(path, tmp_dir, dirs_exist_ok=True, symlinks=True)
_restructure_ui_html_files(tmp_dir)
atexit.register(shutil.rmtree, tmp_dir, True)
verbose_proxy_logger.info(
f"Copied read-only UI to temp dir and restructured: {tmp_dir}"
)
return tmp_dir
except Exception:
shutil.rmtree(tmp_dir, ignore_errors=True)
verbose_proxy_logger.warning(
f"Failed to copy and restructure UI from {path}; "
f"extensionless routes like /ui/login may return 404"
)
return path
try:
ui_path = _prepare_ui_directory(ui_path)
except PermissionError as e:
verbose_proxy_logger.exception(
f"Permission error while restructuring UI directory {ui_path}: {e}"

64
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -646,6 +646,70 @@ def test_read_only_ui_dir_falls_back_to_temp_dir_for_extensionless_routes(tmp_pa
shutil.rmtree(tmp_dir, ignore_errors=True)
def test_prepare_ui_uses_temp_dir_when_not_writable(tmp_path, monkeypatch):
import shutil
from litellm.proxy import proxy_server
ui_root = tmp_path / "ui"
ui_root.mkdir()
(ui_root / "index.html").write_text("<html>index</html>")
(ui_root / "login.html").write_text("<html>login</html>")
(ui_root / "_next").mkdir()
monkeypatch.setattr(proxy_server.os, "access", lambda path, mode: False)
result = proxy_server._prepare_ui_directory(str(ui_root))
try:
assert result != str(ui_root)
assert (Path(result) / "login" / "index.html").exists()
assert not (Path(result) / "login.html").exists()
app = FastAPI()
app.mount("/ui", StaticFiles(directory=result, html=True), name="ui")
response = TestClient(app).get("/ui/login")
assert response.status_code == 200
assert "login" in response.text
finally:
shutil.rmtree(result, ignore_errors=True)
def test_prepare_ui_cleans_up_temp_dir_on_copy_failure(tmp_path, monkeypatch):
import shutil
from litellm.proxy import proxy_server
ui_root = tmp_path / "ui"
ui_root.mkdir()
(ui_root / "index.html").write_text("<html>index</html>")
(ui_root / "_next").mkdir()
monkeypatch.setattr(proxy_server.os, "access", lambda path, mode: False)
monkeypatch.setattr(
proxy_server.shutil,
"copytree",
lambda *a, **kw: (_ for _ in ()).throw(OSError("disk full")),
)
import tempfile
created_dirs: list = []
real_mkdtemp = tempfile.mkdtemp
def tracking_mkdtemp(**kwargs):
d = real_mkdtemp(**kwargs)
created_dirs.append(d)
return d
monkeypatch.setattr(proxy_server.tempfile, "mkdtemp", tracking_mkdtemp)
result = proxy_server._prepare_ui_directory(str(ui_root))
assert result == str(ui_root)
for d in created_dirs:
assert not Path(d).exists(), f"temp dir {d} was not cleaned up"
def test_admin_ui_export_serves_nested_extensionless_routes():
out_dir = Path(litellm.__file__).parent / "proxy" / "_experimental" / "out"
assert out_dir.is_dir(), f"missing UI export at {out_dir}"