From 8338bd9c539aa48d2a1bb20e593b277a46743074 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Mon, 8 Dec 2025 16:30:37 -0800
Subject: [PATCH] Change deprecation banner to only show on /sso/key/generate

---
 .../proxy/common_utils/html_forms/ui_login.py | 26 ++++++++++---
 litellm/proxy/proxy_server.py                 | 10 +++--
 tests/test_litellm/proxy/test_proxy_server.py | 38 +++++++++++++++++++
 3 files changed, 65 insertions(+), 9 deletions(-)
diff --git a/litellm/proxy/common_utils/html_forms/ui_login.py b/litellm/proxy/common_utils/html_forms/ui_login.py
index 8478d41e47..42cfb592a7 100644
--- a/litellm/proxy/common_utils/html_forms/ui_login.py
+++ b/litellm/proxy/common_utils/html_forms/ui_login.py
@@ -8,7 +8,22 @@ if server_root_path != "":
     url_to_redirect_to += server_root_path
 url_to_redirect_to += "/login"
 new_ui_login_url = get_custom_url("", "ui/login")
-html_form = f"""
+
+
+def build_ui_login_form(show_deprecation_banner: bool = False) -> str:
+    banner_html = (
+        f"""
+        <div class="deprecation-banner">
+            <strong>Deprecated:</strong> Logging in with username and password on this page is deprecated.
+            Please use the <a href="{new_ui_login_url}">new login page</a> instead.
+            This page will be dedicated to signing in via SSO in the future.
+        </div>
+        """
+        if show_deprecation_banner
+        else ""
+    )
+
+    return f"""
 <!DOCTYPE html>
 <html lang="en">
 <head>
@@ -209,11 +224,7 @@ html_form = f"""
 </head>
 <body>
     <form action="{url_to_redirect_to}" method="post">
-        <div class="deprecation-banner">
-            <strong>Deprecated:</strong> Logging in with username and password on this page is deprecated.
-            Please use the <a href="{new_ui_login_url}">new login page</a> instead.
-            This page will be dedicated to signing in via SSO in the future.
-        </div>
+        {banner_html}
         <div class="logo-container">
             <div class="logo">
                 🚅 LiteLLM
@@ -253,3 +264,6 @@ html_form = f"""
 </body>
 </html>
 """
+
+
+html_form = build_ui_login_form(show_deprecation_banner=True)
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 5d097ac440..6aef4ab493 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -236,7 +236,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
     decrypt_value_helper,
     encrypt_value_helper,
 )
-from litellm.proxy.common_utils.html_forms.ui_login import html_form
+from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form
 from litellm.proxy.common_utils.http_parsing_utils import (
     _read_request_body,
     check_file_size_under_limit,
@@ -8306,11 +8306,15 @@ async def fallback_login(request: Request):
         # Use UI Credentials set in .env
         from fastapi.responses import HTMLResponse
 
-        return HTMLResponse(content=html_form, status_code=200)
+        return HTMLResponse(
+            content=build_ui_login_form(show_deprecation_banner=False), status_code=200
+        )
     else:
         from fastapi.responses import HTMLResponse
 
-        return HTMLResponse(content=html_form, status_code=200)
+        return HTMLResponse(
+            content=build_ui_login_form(show_deprecation_banner=False), status_code=200
+        )
 
 
 @router.post(
diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py
index 76fd37596e..7b99cd0b92 100644
--- a/tests/test_litellm/proxy/test_proxy_server.py
+++ b/tests/test_litellm/proxy/test_proxy_server.py
@@ -124,6 +124,44 @@ def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch):
     )
 
 
+def test_fallback_login_has_no_deprecation_banner(client_no_auth):
+    response = client_no_auth.get("/fallback/login")
+
+    assert response.status_code == 200
+    html = response.text
+    assert '<div class="deprecation-banner">' not in html
+    assert "Deprecated:" not in html
+    assert "<form" in html
+
+
+def test_sso_key_generate_shows_deprecation_banner(client_no_auth, monkeypatch):
+    # Ensure the route returns the HTML form instead of redirecting
+    monkeypatch.setattr(
+        "litellm.proxy.management_endpoints.ui_sso.show_missing_vars_in_env",
+        lambda: None,
+    )
+    monkeypatch.setattr(
+        "litellm.proxy.management_endpoints.ui_sso.SSOAuthenticationHandler.get_redirect_url_for_sso",
+        lambda *args, **kwargs: "http://test/redirect",
+    )
+    monkeypatch.setattr(
+        "litellm.proxy.management_endpoints.ui_sso.SSOAuthenticationHandler._get_cli_state",
+        lambda *args, **kwargs: None,
+    )
+    monkeypatch.setattr(
+        "litellm.proxy.management_endpoints.ui_sso.SSOAuthenticationHandler.should_use_sso_handler",
+        lambda *args, **kwargs: False,
+    )
+    monkeypatch.setenv("UI_USERNAME", "admin")
+
+    response = client_no_auth.get("/sso/key/generate")
+
+    assert response.status_code == 200
+    html = response.text
+    assert '<div class="deprecation-banner">' in html
+    assert "Deprecated:" in html
+
+
 @pytest.mark.asyncio
 async def test_initialize_scheduled_jobs_credentials(monkeypatch):
     """
