diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 2dcb7cb4787a..bb656e04e5b1 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -36,6 +36,7 @@ RUN cd /app/ui/litellm-dashboard && npm install --legacy-peer-deps RUN cd /app/ui/litellm-dashboard && npm run build RUN cp -r /app/ui/litellm-dashboard/out/* /tmp/litellm_ui/ +RUN mkdir -p /tmp/litellm_assets && cp /app/litellm/proxy/logo.jpg /tmp/litellm_assets/logo.jpg RUN cd /tmp/litellm_ui && \ for html_file in *.html; do \ @@ -72,6 +73,7 @@ COPY --from=builder /app/schema.prisma /app/schema.prisma COPY --from=builder /app/dist/*.whl . COPY --from=builder /wheels/ /wheels/ COPY --from=builder /tmp/litellm_ui /tmp/litellm_ui +COPY --from=builder /tmp/litellm_assets /tmp/litellm_assets # Install package from wheel and dependencies RUN pip install *.whl /wheels/* --no-index --find-links=/wheels/ \ @@ -100,8 +102,8 @@ RUN pip install --no-cache-dir prisma && \ chmod +x docker/prod_entrypoint.sh # Create directories and set permissions for non-root user -RUN mkdir -p /nonexistent /.npm && \ - chown -R nobody:nogroup /app /tmp/litellm_ui /nonexistent /.npm && \ +RUN mkdir -p /nonexistent /.npm /tmp/litellm_assets && \ + chown -R nobody:nogroup /app /tmp/litellm_ui /tmp/litellm_assets /nonexistent /.npm && \ PRISMA_PATH=$(python -c "import os, prisma; print(os.path.dirname(prisma.__file__))") && \ chown -R nobody:nogroup $PRISMA_PATH && \ LITELLM_PKG_MIGRATIONS_PATH="$(python -c 'import os, litellm_proxy_extras; print(os.path.dirname(litellm_proxy_extras.__file__))' 2>/dev/null || echo '')/migrations" && \ @@ -110,11 +112,11 @@ RUN mkdir -p /nonexistent /.npm && \ # OpenShift compatibility RUN PRISMA_PATH=$(python -c "import os, prisma; print(os.path.dirname(prisma.__file__))") && \ LITELLM_PROXY_EXTRAS_PATH=$(python -c "import os, litellm_proxy_extras; print(os.path.dirname(litellm_proxy_extras.__file__))" 2>/dev/null || echo "") && \ - chgrp -R 0 $PRISMA_PATH /tmp/litellm_ui && \ + chgrp -R 0 $PRISMA_PATH /tmp/litellm_ui /tmp/litellm_assets && \ [ -n "$LITELLM_PROXY_EXTRAS_PATH" ] && chgrp -R 0 $LITELLM_PROXY_EXTRAS_PATH || true && \ - chmod -R g=u $PRISMA_PATH /tmp/litellm_ui && \ + chmod -R g=u $PRISMA_PATH /tmp/litellm_ui /tmp/litellm_assets && \ [ -n "$LITELLM_PROXY_EXTRAS_PATH" ] && chmod -R g=u $LITELLM_PROXY_EXTRAS_PATH || true && \ - chmod -R g+w $PRISMA_PATH /tmp/litellm_ui && \ + chmod -R g+w $PRISMA_PATH /tmp/litellm_ui /tmp/litellm_assets && \ [ -n "$LITELLM_PROXY_EXTRAS_PATH" ] && chmod -R g+w $LITELLM_PROXY_EXTRAS_PATH || true # Switch to non-root user diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 28ee370ff0c7..57c2cff8352f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -8720,7 +8720,19 @@ def get_image(): # get current_dir current_dir = os.path.dirname(os.path.abspath(__file__)) - default_logo = os.path.join(current_dir, "logo.jpg") + default_site_logo = os.path.join(current_dir, "logo.jpg") + + is_non_root = os.getenv("LITELLM_NON_ROOT", "").lower() == "true" + assets_dir = "/tmp/litellm_assets" if is_non_root else current_dir + + if is_non_root: + os.makedirs(assets_dir, exist_ok=True) + + default_logo = ( + os.path.join(assets_dir, "logo.jpg") if is_non_root else default_site_logo + ) + if is_non_root and not os.path.exists(default_logo): + default_logo = default_site_logo logo_path = os.getenv("UI_LOGO_PATH", default_logo) verbose_proxy_logger.debug("Reading logo from path: %s", logo_path) @@ -8732,7 +8744,8 @@ def get_image(): response = client.get(logo_path) if response.status_code == 200: # Save the image to a local file - cache_path = os.path.join(current_dir, "cached_logo.jpg") + cache_dir = assets_dir if is_non_root else current_dir + cache_path = os.path.join(cache_dir, "cached_logo.jpg") with open(cache_path, "wb") as f: f.write(response.content) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 8e87b67933fc..ccbf974b725c 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -2494,3 +2494,134 @@ def test_get_prompt_spec_for_db_prompt_with_versions(): prompt_spec_v2 = proxy_config._get_prompt_spec_for_db_prompt(db_prompt=mock_prompt_v2) assert prompt_spec_v2.prompt_id == "chat_prompt.v2" + +def test_get_image_non_root_uses_tmp_assets_dir(monkeypatch): + """ + Test that get_image uses /tmp/litellm_assets when LITELLM_NON_ROOT is true. + """ + from unittest.mock import patch + + from litellm.proxy.proxy_server import get_image + + # Set LITELLM_NON_ROOT to true + monkeypatch.setenv("LITELLM_NON_ROOT", "true") + monkeypatch.delenv("UI_LOGO_PATH", raising=False) + + # Mock os.path operations + with patch("litellm.proxy.proxy_server.os.makedirs") as mock_makedirs, \ + patch("litellm.proxy.proxy_server.os.path.exists", return_value=True), \ + patch("litellm.proxy.proxy_server.os.getenv") as mock_getenv, \ + patch("litellm.proxy.proxy_server.FileResponse") as mock_file_response: + + # Setup mock_getenv to return empty string for UI_LOGO_PATH + def getenv_side_effect(key, default=""): + if key == "UI_LOGO_PATH": + return "" + elif key == "LITELLM_NON_ROOT": + return "true" + return default + + mock_getenv.side_effect = getenv_side_effect + + # Call the function + get_image() + + # Verify makedirs was called with /tmp/litellm_assets + mock_makedirs.assert_called_once_with("/tmp/litellm_assets", exist_ok=True) + + +def test_get_image_non_root_fallback_to_default_logo(monkeypatch): + """ + Test that get_image falls back to default_site_logo when logo doesn't exist + in /tmp/litellm_assets for non-root case. + """ + from unittest.mock import patch + + from litellm.proxy.proxy_server import get_image + + # Set LITELLM_NON_ROOT to true + monkeypatch.setenv("LITELLM_NON_ROOT", "true") + monkeypatch.delenv("UI_LOGO_PATH", raising=False) + + # Track path.exists calls to verify it checks /tmp/litellm_assets/logo.jpg + exists_calls = [] + + def exists_side_effect(path): + exists_calls.append(path) + # Return False for /tmp/litellm_assets/logo.jpg to trigger fallback + if "/tmp/litellm_assets/logo.jpg" in path: + return False + return True + + # Mock os.path operations + with patch("litellm.proxy.proxy_server.os.makedirs") as mock_makedirs, \ + patch("litellm.proxy.proxy_server.os.path.exists", side_effect=exists_side_effect), \ + patch("litellm.proxy.proxy_server.os.getenv") as mock_getenv, \ + patch("litellm.proxy.proxy_server.FileResponse") as mock_file_response: + + # Setup mock_getenv + def getenv_side_effect(key, default=""): + if key == "UI_LOGO_PATH": + return "" + elif key == "LITELLM_NON_ROOT": + return "true" + return default + + mock_getenv.side_effect = getenv_side_effect + + # Call the function + get_image() + + # Verify makedirs was called with /tmp/litellm_assets + mock_makedirs.assert_called_once_with("/tmp/litellm_assets", exist_ok=True) + + # Verify that exists was called to check /tmp/litellm_assets/logo.jpg + tmp_logo_path = "/tmp/litellm_assets/logo.jpg" + assert any(tmp_logo_path in str(call) for call in exists_calls), \ + f"Should check if {tmp_logo_path} exists" + + # Verify FileResponse was called (with fallback logo) + assert mock_file_response.called, "FileResponse should be called" + + +def test_get_image_root_case_uses_current_dir(monkeypatch): + """ + Test that get_image uses current_dir when LITELLM_NON_ROOT is not true. + """ + from unittest.mock import patch + + from litellm.proxy.proxy_server import get_image + + # Don't set LITELLM_NON_ROOT (or set it to false) + monkeypatch.delenv("LITELLM_NON_ROOT", raising=False) + monkeypatch.delenv("UI_LOGO_PATH", raising=False) + + # Mock os.path operations + with patch("litellm.proxy.proxy_server.os.makedirs") as mock_makedirs, \ + patch("litellm.proxy.proxy_server.os.path.exists", return_value=True), \ + patch("litellm.proxy.proxy_server.os.getenv") as mock_getenv, \ + patch("litellm.proxy.proxy_server.FileResponse") as mock_file_response: + + # Setup mock_getenv + def getenv_side_effect(key, default=""): + if key == "UI_LOGO_PATH": + return "" + elif key == "LITELLM_NON_ROOT": + return "" # Not set or empty + return default + + mock_getenv.side_effect = getenv_side_effect + + # Call the function + get_image() + + # Verify makedirs was NOT called with /tmp/litellm_assets (should not create it for root case) + tmp_assets_calls = [ + call for call in mock_makedirs.call_args_list + if "/tmp/litellm_assets" in str(call) + ] + assert len(tmp_assets_calls) == 0, "Should not create /tmp/litellm_assets for root case" + + # Verify FileResponse was called + assert mock_file_response.called, "FileResponse should be called" +