From 82a7bf12a6e15be5ab956bfdbbd9c16765641cce Mon Sep 17 00:00:00 2001 From: Daniele Salvador Date: Mon, 1 Jun 2026 01:05:10 +0200 Subject: [PATCH] refactor(proxy): extract _prepare_ui_directory; add symlinks=True; cover read-only branch directly 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. --- litellm/proxy/proxy_server.py | 55 ++++++++-------- tests/test_litellm/proxy/test_proxy_server.py | 64 +++++++++++++++++++ 2 files changed, 94 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bad9b807f34..112bbaacc93 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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}" diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 5b2b07cfeea..28fffc237cf 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -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("index") + (ui_root / "login.html").write_text("login") + (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("index") + (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}"