Remove added tunnel domains to Customer Application on process exit (#…

…1966)

* Copy .shopify back to diff directory

* Return cleanup function from CA push

* Ensure codegen process doesn't print on normal exit

* Ensure MiniOxygen is disposed when Vite closes

* Cleanup resources on process exit

* Doc

packages/cli/src/commands/hydrogen/customer-account/push.ts

@@ -1,16 +1,17 @@
 import Command from '@shopify/cli-kit/node/base-command';
 import {Flags} from '@oclif/core';
 import {AbortError} from '@shopify/cli-kit/node/error';
+import {outputDebug} from '@shopify/cli-kit/node/output';
 import {linkStorefront} from '../link.js';
 import {commonFlags, flagsToCamelObject} from '../../../lib/flags.js';
 import {getCliCommand} from '../../../lib/shell.js';
-import {login} from '../../../lib/auth.js';
+import {type AdminSession, login} from '../../../lib/auth.js';
 import {
+  type ShopifyConfig,
   getConfig,
   setCustomerAccountConfig,
 } from '../../../lib/shopify-config.js';
 import {replaceCustomerApplicationUrls} from '../../../lib/graphql/admin/customer-application-update.js';
-import {FatalErrorType} from '@shopify/cli-kit/node/error';
 
 export default class CustomerAccountPush extends Command {
   static description = 'Push project configuration to admin';
@@ -110,6 +111,13 @@ export async function runCustomerAccountPush({
       javascriptOrigin,
       logoutUri,
     });
+
+    return () =>
+      cleanupCustomerApplicationUrls(session, storefrontId, {
+        redirectUri,
+        javascriptOrigin,
+        logoutUri,
+      });
   } catch (error: any) {
     let confidentialAccessFound = false;
 
@@ -167,6 +175,33 @@ export async function runCustomerAccountPush({
   }
 }
 
+/**
+ * Removes existing Customer Application URLs from the storefront.
+ */
+async function cleanupCustomerApplicationUrls(
+  session: AdminSession,
+  storefrontId: string,
+  customerAccountConfig: NonNullable<
+    NonNullable<ShopifyConfig['storefront']>['customerAccountConfig']
+  > = {},
+) {
+  if (!Object.values(customerAccountConfig).some(Boolean)) return;
+
+  outputDebug(
+    `Cleaning up Customer Application url "${customerAccountConfig.redirectUri}"`,
+  );
+
+  await replaceCustomerApplicationUrls(session, storefrontId, {
+    redirectUri: {removeRegex: customerAccountConfig?.redirectUri},
+    javascriptOrigin: {removeRegex: customerAccountConfig?.javascriptOrigin},
+    logoutUris: {removeRegex: customerAccountConfig?.logoutUri},
+  }).catch((error) => {
+    outputDebug(
+      `Failed to clean up Customer Application url "${customerAccountConfig.redirectUri}":\n${error?.message}`,
+    );
+  });
+}
+
 export async function getStorefrontId(
   root: string,
   storefrontIdFromFlag?: string,

packages/cli/src/commands/hydrogen/dev-vite.ts

@@ -215,14 +215,6 @@ export async function runViteDev({
     ],
   });
 
-  process.once('SIGTERM', async () => {
-    try {
-      await viteServer.close();
-    } finally {
-      process.exit();
-    }
-  });
-
   const h2Plugin = findHydrogenPlugin(viteServer.config);
   if (!h2Plugin) {
     await viteServer.close();
@@ -258,7 +250,7 @@ export async function runViteDev({
   //   process.env.HYDROGEN_ASSET_BASE_URL = buildAssetsUrl(assetsPort);
   // }
 
-  const [tunnelHost, cliCommand] = await Promise.all([
+  const [tunnel, cliCommand] = await Promise.all([
     backgroundPromise.then(({customerAccountPush, storefrontId}) =>
       customerAccountPush
         ? startTunnelAndPushConfig(root, cliConfig, publicPort, storefrontId)
@@ -272,7 +264,7 @@ export async function runViteDev({
     viteServer.resolvedUrls!.local[0] ?? viteServer.resolvedUrls!.network[0]!,
   );
 
-  const finalHost = tunnelHost || publicUrl.toString() || publicUrl.origin;
+  const finalHost = tunnel?.host || publicUrl.toString() || publicUrl.origin;
 
   // Start the public facing server with the port passed by the user.
   enhanceH2Logs({
@@ -323,8 +315,9 @@ export async function runViteDev({
   return {
     getUrl: () => finalHost,
     async close() {
-      codegenProcess?.kill(0);
-      await viteServer.close();
+      codegenProcess?.removeAllListeners('close');
+      codegenProcess?.kill('SIGINT');
+      await Promise.allSettled([viteServer.close(), tunnel?.cleanup?.()]);
     },
   };
 }

packages/cli/src/commands/hydrogen/dev.ts

@@ -1,9 +1,9 @@
-import path from 'node:path';
 import fs from 'node:fs/promises';
 import type {ChildProcess} from 'node:child_process';
 import {outputDebug, outputInfo} from '@shopify/cli-kit/node/output';
 import {fileExists} from '@shopify/cli-kit/node/fs';
 import {renderFatalError} from '@shopify/cli-kit/node/ui';
+import {relativePath, resolvePath} from '@shopify/cli-kit/node/path';
 import {copyPublicFiles} from './build.js';
 import {
   type RemixConfig,
@@ -40,7 +40,10 @@ import {setupLiveReload} from '../../lib/live-reload.js';
 import {checkRemixVersions} from '../../lib/remix-version-check.js';
 import {displayDevUpgradeNotice} from './upgrade.js';
 import {findPort} from '../../lib/find-port.js';
-import {prepareDiffDirectory} from '../../lib/template-diff.js';
+import {
+  copyShopifyConfig,
+  prepareDiffDirectory,
+} from '../../lib/template-diff.js';
 import {
   startTunnelAndPushConfig,
   getDevConfigInBackground,
@@ -89,7 +92,10 @@ export default class Dev extends Command {
 
   async run(): Promise<void> {
     const {flags} = await this.parse(Dev);
-    let directory = flags.path ? path.resolve(flags.path) : process.cwd();
+    const originalDirectory = flags.path
+      ? resolvePath(flags.path)
+      : process.cwd();
+    let directory = originalDirectory;
 
     if (flags.diff) {
       directory = await prepareDiffDirectory(directory, true);
@@ -102,11 +108,29 @@ export default class Dev extends Command {
       cliConfig: this.config,
     };
 
-    if (await hasViteConfig(directory ?? process.cwd())) {
-      const {runViteDev} = await import('./dev-vite.js');
-      await runViteDev(devParams);
-    } else {
-      await runDev(devParams);
+    const {close} = (await hasViteConfig(directory ?? process.cwd()))
+      ? await import('./dev-vite.js').then(({runViteDev}) =>
+          runViteDev(devParams),
+        )
+      : await runDev(devParams);
+
+    // Note: Shopify CLI is hooking into process events and calling process.exit.
+    // This means we are unable to hook into 'beforeExit' or 'SIGINT" events
+    // to cleanup resources. In addition, Miniflare uses `exit-hook` dependency
+    // to do the same thing. This is a workaround to ensure we cleanup resources:
+    let closingPromise: Promise<void>;
+    const processExit = process.exit;
+    // @ts-expect-error - Async function
+    process.exit = async (code?: number | undefined) => {
+      // This function will be called multiple times,
+      // but we only want to cleanup resources once.
+      closingPromise ??= close();
+      await closingPromise;
+      return processExit(code);
+    };
+
+    if (flags.diff) {
+      await copyShopifyConfig(directory, originalDirectory);
     }
   }
 }
@@ -178,8 +202,8 @@ export async function runDev({
   };
 
   const getFilePaths = (file: string) => {
-    const fileRelative = path.relative(root, file);
-    return [fileRelative, path.resolve(root, fileRelative)] as const;
+    const fileRelative = relativePath(root, file);
+    return [fileRelative, resolvePath(root, fileRelative)] as const;
   };
 
   const serverBundleExists = () => fileExists(buildPathWorkerFile);
@@ -263,7 +287,7 @@ export async function runDev({
 
     logInjectedVariables();
 
-    const host = (await tunnelPromise) ?? miniOxygen.listeningAt;
+    const host = (await tunnelPromise)?.host ?? miniOxygen.listeningAt;
 
     const cliCommand = await cliCommandPromise;
     enhanceH2Logs({host, cliCommand, ...remixConfig});
@@ -412,8 +436,13 @@ export async function runDev({
   return {
     getUrl: () => miniOxygen.listeningAt,
     async close() {
-      codegenProcess?.kill(0);
-      await Promise.all([closeWatcher(), miniOxygen?.close()]);
+      codegenProcess?.removeAllListeners('close');
+      codegenProcess?.kill('SIGINT');
+      await Promise.allSettled([
+        closeWatcher(),
+        miniOxygen?.close(),
+        Promise.resolve(tunnelPromise).then((tunnel) => tunnel?.cleanup?.()),
+      ]);
     },
   };
 }

packages/cli/src/lib/dev-shared.ts

@@ -84,23 +84,21 @@ export async function startTunnelAndPushConfig(
     host.replace(TUNNEL_DOMAIN.ORIGINAL, TUNNEL_DOMAIN.REBRANDED),
   );
 
-  try {
-    await runCustomerAccountPush({
-      path: root,
-      devOrigin: host,
-      storefrontId,
-    });
-  } catch (error) {
+  const cleanup = await runCustomerAccountPush({
+    path: root,
+    devOrigin: host,
+    storefrontId,
+  }).catch((error) => {
     if (error instanceof AbortError) {
       renderInfo({
         headline: 'Customer Account Application setup update fail.',
         body: error.tryMessage || undefined,
         nextSteps: error.nextSteps,
       });
     }
-  }
+  });
 
-  return host;
+  return {host, cleanup};
 }
 
 export function getDebugBannerLine(publicInspectorPort: number) {

packages/cli/src/lib/template-diff.ts

@@ -1,7 +1,11 @@
 import {rmdirSync} from 'node:fs';
 import {temporaryDirectory} from 'tempy';
 import {createSymlink, copy as copyDirectory} from 'fs-extra/esm';
-import {copyFile, removeFile as remove} from '@shopify/cli-kit/node/fs';
+import {
+  copyFile,
+  fileExists,
+  removeFile as remove,
+} from '@shopify/cli-kit/node/fs';
 import {joinPath, relativePath} from '@shopify/cli-kit/node/path';
 import {readAndParsePackageJson} from '@shopify/cli-kit/node/node-package-manager';
 import colors from '@shopify/cli-kit/node/colors';
@@ -145,19 +149,42 @@ export async function applyTemplateDiff(
   });
 }
 
+/**
+ * Brings the `dist` directory back to the original project.
+ * This is used to run `h2 preview` with the resulting build.
+ */
 export async function copyDiffBuild(
-  targetDirectory: string,
+  generatedDirectory: string,
   diffDirectory: string,
 ) {
-  const targetDist = joinPath(diffDirectory, 'dist');
-  await remove(targetDist);
+  const target = joinPath(diffDirectory, 'dist');
+  await remove(target);
   await Promise.all([
-    copyDirectory(joinPath(targetDirectory, 'dist'), targetDist, {
+    copyDirectory(joinPath(generatedDirectory, 'dist'), target, {
       overwrite: true,
     }),
     copyFile(
-      joinPath(targetDirectory, '.env'),
+      joinPath(generatedDirectory, '.env'),
       joinPath(diffDirectory, '.env'),
     ),
   ]);
 }
+
+/**
+ * Brings the `.shopify` directory back to the original project.
+ * This is important to keep a reference of the tunnel configuration
+ * so that it can be removed in the next run.
+ */
+export async function copyShopifyConfig(
+  generatedDirectory: string,
+  diffDirectory: string,
+) {
+  const source = joinPath(generatedDirectory, '.shopify');
+  if (!(await fileExists(source))) return;
+
+  const target = joinPath(diffDirectory, '.shopify');
+  await remove(target);
+  await copyDirectory(joinPath(generatedDirectory, '.shopify'), target, {
+    overwrite: true,
+  });
+}

packages/mini-oxygen/src/vite/server-middleware.ts

@@ -162,6 +162,12 @@ export async function startMiniOxygenRuntime({
     ? warmupWorkerdCache()
     : viteDevServer.httpServer?.once('listening', warmupWorkerdCache);
 
+  // Ensure MiniOxygen is disposed when Vite is closed
+  const viteClose = viteDevServer.close;
+  viteDevServer.close = async () => {
+    await Promise.allSettled([viteClose(), miniOxygen.dispose()]);
+  };
+
   return miniOxygen;
 }