diff --git a/edsdme/scripts/utils.js b/edsdme/scripts/utils.js
index bbb811c..a70ec57 100644
--- a/edsdme/scripts/utils.js
+++ b/edsdme/scripts/utils.js
@@ -29,11 +29,15 @@ export const [setLibs, getLibs] = (() => {
   return [
     (prodLibs, location) => {
       libs = (() => {
-        const { hostname, search } = location || window.location;
-        // TODO: check if better ways are possible for partners.stage.adobe.com
-        if (!(hostname.includes('.hlx.') || hostname.includes('local') || hostname === 'partners.stage.adobe.com' || hostname === 'partners.adobe.com')) return prodLibs;
-        const branch = new URLSearchParams(search).get('milolibs') || 'main';
-        if (branch === 'local') return 'http://localhost:6456/libs';
+        const { hostname, search, origin } = location || window.location;
+        if (origin.endsWith('adobe.com')) {
+          return origin.replace('partners', 'milo') + prodLibs;
+        }
+        const partnerBranch = hostname.startsWith('main') ? 'main' : 'stage';
+        const branch = new URLSearchParams(search).get('milolibs') || partnerBranch;
+        if (branch === 'local') {
+          return 'http://localhost:6456/libs';
+        }
         return branch.includes('--') ? `https://${branch}.hlx.live/libs` : `https://${branch}--milo--adobecom.hlx.live/libs`;
       })();
       return libs;
diff --git a/test/scripts/utils.jest.js b/test/scripts/utils.jest.js
index 71b7d53..cd9c5c4 100644
--- a/test/scripts/utils.jest.js
+++ b/test/scripts/utils.jest.js
@@ -52,7 +52,7 @@ describe('Test utils.js', () => {
   it('Milo libs', () => {
     window.location.hostname = 'partners.stage.adobe.com';
     const libs = setLibs('/libs');
-    expect(libs).toEqual('https://main--milo--adobecom.hlx.live/libs');
+    expect(libs).toEqual('https://milo.stage.adobe.com/libs');
   });
   describe('Test update footer and gnav', () => {
     beforeEach(() => {
diff --git a/test/scripts/utils.test.js b/test/scripts/utils.test.js
index 41c52d0..903ec7e 100644
--- a/test/scripts/utils.test.js
+++ b/test/scripts/utils.test.js
@@ -4,21 +4,40 @@ import { setLibs } from '../../edsdme/scripts/utils.js';
 describe('Libs', () => {
   it('Default Libs', () => {
     const libs = setLibs('/libs');
+    expect(libs).to.equal('https://stage--milo--adobecom.hlx.live/libs');
+  });
+  it('Main Libs', () => {
+    const location = {
+      hostname: 'main--dme-partners.hlx.page',
+      origin: 'https://main--dme-partners.hlx.page',
+    };
+    const libs = setLibs('/libs', location);
     expect(libs).to.equal('https://main--milo--adobecom.hlx.live/libs');
   });
+  it('Returns prod milo for prod', () => {
+    const location = { origin: 'https://partners.adobe.com' };
+    const libs = setLibs('/libs', location);
+    expect(libs).to.equal('https://milo.adobe.com/libs');
+  });
+  it('Returns stage milo for stage', () => {
+    const location = { origin: 'https://partners.stage.adobe.com' };
+    const libs = setLibs('/libs', location);
+    expect(libs).to.equal('https://milo.stage.adobe.com/libs');
+  });
 
   it('Does not support milolibs query param on prod', () => {
     const location = {
-      hostname: 'business.adobe.com',
+      origin: 'https://partners.adobe.com',
       search: '?milolibs=foo',
     };
     const libs = setLibs('/libs', location);
-    expect(libs).to.equal('/libs');
+    expect(libs).to.equal('https://milo.adobe.com/libs');
   });
 
   it('Supports milolibs query param', () => {
     const location = {
       hostname: 'localhost',
+      origin: 'http://localhost:3000',
       search: '?milolibs=foo',
     };
     const libs = setLibs('/libs', location);
@@ -28,6 +47,7 @@ describe('Libs', () => {
   it('Supports local milolibs query param', () => {
     const location = {
       hostname: 'localhost',
+      origin: 'http://localhost:3000',
       search: '?milolibs=local',
     };
     const libs = setLibs('/libs', location);
@@ -37,6 +57,7 @@ describe('Libs', () => {
   it('Supports forked milolibs query param', () => {
     const location = {
       hostname: 'localhost',
+      origin: 'http://localhost:3000',
       search: '?milolibs=awesome--milo--forkedowner',
     };
     const libs = setLibs('/libs', location);