diff --git a/code/addons/viewport/src/Tool.tsx b/code/addons/viewport/src/Tool.tsx
index 9a0d6c04f49a..df65988e5097 100644
--- a/code/addons/viewport/src/Tool.tsx
+++ b/code/addons/viewport/src/Tool.tsx
@@ -106,9 +106,10 @@ const getStyles = (
   styles: Styles,
   isRotated: boolean
 ): ViewportStyles | undefined => {
-  if (!styles || !prevStyles) {
+  if (styles === null) {
     return undefined;
   }
+
   const result = typeof styles === 'function' ? styles(prevStyles) : styles;
   return isRotated ? flip(result) : result;
 };
diff --git a/code/addons/viewport/src/models/Viewport.ts b/code/addons/viewport/src/models/Viewport.ts
index 76a5eb096777..7b3c98e069c9 100644
--- a/code/addons/viewport/src/models/Viewport.ts
+++ b/code/addons/viewport/src/models/Viewport.ts
@@ -1,4 +1,4 @@
-export type Styles = ViewportStyles | ((s: ViewportStyles) => ViewportStyles) | null;
+export type Styles = ViewportStyles | ((s: ViewportStyles | undefined) => ViewportStyles) | null;
 
 export interface Viewport {
   name: string;
diff --git a/code/e2e-tests/addon-viewport.spec.ts b/code/e2e-tests/addon-viewport.spec.ts
index 18d72e112646..da66778a3708 100644
--- a/code/e2e-tests/addon-viewport.spec.ts
+++ b/code/e2e-tests/addon-viewport.spec.ts
@@ -20,4 +20,24 @@ test.describe('addon-viewport', () => {
     // Check that Button story is still displayed
     await expect(sbPage.previewRoot()).toContainText('Button');
   });
+
+  test('iframe width should be changed when a mobile viewport is selected', async ({ page }) => {
+    const sbPage = new SbPage(page);
+
+    // Click on viewport button and select small mobile
+    await sbPage.navigateToStory('example/button', 'primary');
+
+    // Measure the original dimensions of previewRoot
+    const originalDimensions = await sbPage.previewRoot().boundingBox();
+    await expect(originalDimensions?.width).toBeDefined();
+
+    await sbPage.selectToolbar('[title="Change the size of the preview"]', '#list-item-mobile1');
+
+    // Measure the adjusted dimensions of previewRoot after clicking the mobile item.
+    const adjustedDimensions = await sbPage.previewRoot().boundingBox();
+    await expect(adjustedDimensions?.width).toBeDefined();
+
+    // Compare the two widths
+    await expect(adjustedDimensions?.width).not.toBe(originalDimensions?.width);
+  });
 });