diff --git a/code/lib/builder-vite/src/codegen-iframe-script.ts b/code/lib/builder-vite/src/codegen-iframe-script.ts
index cc35bb9071ee..5dc0e792655d 100644
--- a/code/lib/builder-vite/src/codegen-iframe-script.ts
+++ b/code/lib/builder-vite/src/codegen-iframe-script.ts
@@ -39,6 +39,9 @@ export async function generateIframeScriptCode(options: ExtendedOptions) {
       addDecorator,
       addParameters,
       addLoader,
+      addArgs,
+      addArgTypes,
+      addStepRunner,
       addArgTypesEnhancer,
       addArgsEnhancer,
       setGlobalRender,
@@ -53,22 +56,10 @@ export async function generateIframeScriptCode(options: ExtendedOptions) {
         const value = config[key];
         switch (key) {
           case 'args': {
-            if (typeof clientApi.addArgs !== "undefined") {
-              return clientApi.addArgs(value);
-            } else {
-              return logger.warn(
-                "Could not add global args. Please open an issue in storybookjs/builder-vite."
-              );
-            }
+            return addArgs(value);
           }
           case 'argTypes': {
-            if (typeof clientApi.addArgTypes !== "undefined") {
-              return clientApi.addArgTypes(value);
-            } else {
-              return logger.warn(
-                "Could not add global argTypes. Please open an issue in storybookjs/builder-vite."
-              );
-            }
+            return addArgTypes(value);
           }
           case 'decorators': {
             return value.forEach((decorator) => addDecorator(decorator, false));
@@ -99,6 +90,9 @@ export async function generateIframeScriptCode(options: ExtendedOptions) {
           case 'renderToDOM': {
             return null; // This key is not handled directly in v6 mode.
           }
+          case 'runStep': {
+            return addStepRunner(value);
+          }
           default: {
             // eslint-disable-next-line prefer-template
             return console.log(key + ' was not supported :( !');