diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 6e2e2bedac1..ff3bbf47389 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -1798,7 +1798,10 @@ async def cli_sso_callback( from fastapi.responses import HTMLResponse - verify_url = str(request.url_for("cli_sso_complete", login_id=key)) + verify_url = get_custom_url( + request_base_url=str(request.base_url), + route=f"sso/cli/complete/{key}", + ) html_content = _render_cli_sso_verification_page( verify_url=verify_url, browser_complete_token=browser_complete_token, diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index 83317157847..23216542f35 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -2218,6 +2218,7 @@ class TestCLIKeyRegenerationFlow: # Mock request mock_request = MagicMock(spec=Request) + mock_request.base_url = "http://internal-proxy.local/" # Test data session_key = "cli-session-4567890" @@ -2242,11 +2243,14 @@ class TestCLIKeyRegenerationFlow: "user_code_verified": False, "session_data": None, } - mock_request.url_for.return_value = ( - "https://test.example.com/sso/cli/complete/cli-session-4567890" - ) - with ( + patch.dict( + os.environ, + { + "PROXY_BASE_URL": "https://test.example.com", + "SERVER_ROOT_PATH": "", + }, + ), patch( "litellm.proxy.management_endpoints.ui_sso.get_user_info_from_db", return_value=mock_user_info, @@ -2290,6 +2294,10 @@ class TestCLIKeyRegenerationFlow: assert result.status_code == 200 # Verify response contains success message (response is HTML) assert result.body is not None + assert ( + 'action="https://test.example.com/sso/cli/complete/cli-session-4567890"' + in result.body.decode() + ) @pytest.mark.asyncio async def test_cli_poll_key_returns_teams_for_selection(self):