Preserve store URL in context with AsyncLocalStorage

Previously, commands that accessed local storage
values could potentially receive an incorrect store
URL when running multiple environments concurrently
due to file overwrites.

This commit wraps each command call in a theme store
context provider, eliminating the need to grab
from the local storage file

Author: dejmedus

.changeset/calm-ears-add.md

@@ -0,0 +1,5 @@
+---
+'@shopify/theme': minor
+---
+
+Ensure commands run in multiple environments use the correct store URL

packages/theme/src/cli/services/local-storage.test.ts

@@ -10,6 +10,8 @@ import {
   removeStorefrontPassword,
   ThemeLocalStorageSchema,
   setThemeStore,
+  getThemeStore,
+  useThemeStoreContext,
 } from './local-storage.js'
 import {inTemporaryDirectory} from '@shopify/cli-kit/node/fs'
 import {LocalStorage} from '@shopify/cli-kit/node/local-storage'
@@ -55,4 +57,55 @@ describe('local-storage', () => {
       })
     })
   })
+
+  describe('getThemeStore', () => {
+    test('selects store from context when inside useThemeStoreContext', async () => {
+      await inTemporaryDirectory(async (cwd) => {
+        const storage = new LocalStorage<ThemeLocalStorageSchema>({cwd})
+        setThemeStore('storage-store.myshopify.com', storage)
+
+        const initialStore = getThemeStore(storage)
+        let insideContextStore: string | undefined
+
+        await useThemeStoreContext('context-store.myshopify.com', async () => {
+          insideContextStore = getThemeStore(storage)
+        })
+
+        const outsideContextStore = getThemeStore(storage)
+
+        expect(initialStore).toBe('storage-store.myshopify.com')
+        expect(outsideContextStore).toBe('storage-store.myshopify.com')
+        expect(insideContextStore).toBe('context-store.myshopify.com')
+      })
+    })
+
+    test('ensures concurrently run commands maintain their own store value', async () => {
+      await inTemporaryDirectory(async (cwd) => {
+        const storage = new LocalStorage<ThemeLocalStorageSchema>({cwd})
+        setThemeStore('storage-store.myshopify.com', storage)
+
+        const results: {[key: string]: string | undefined} = {}
+
+        await Promise.all([
+          useThemeStoreContext('store1.myshopify.com', async () => {
+            await new Promise((resolve) => setTimeout(resolve, 10))
+            results.env1 = getThemeStore(storage)
+          }),
+          useThemeStoreContext('store2.myshopify.com', async () => {
+            await new Promise((resolve) => setTimeout(resolve, 5))
+            results.env2 = getThemeStore(storage)
+          }),
+          useThemeStoreContext('store3.myshopify.com', async () => {
+            results.env3 = getThemeStore(storage)
+          }),
+          (results.env4 = getThemeStore(storage)),
+        ])
+
+        expect(results.env1).toBe('store1.myshopify.com')
+        expect(results.env2).toBe('store2.myshopify.com')
+        expect(results.env3).toBe('store3.myshopify.com')
+        expect(results.env4).toBe('storage-store.myshopify.com')
+      })
+    })
+  })
 })

packages/theme/src/cli/services/local-storage.ts

@@ -1,6 +1,7 @@
 import {BugError} from '@shopify/cli-kit/node/error'
 import {LocalStorage} from '@shopify/cli-kit/node/local-storage'
 import {outputDebug, outputContent} from '@shopify/cli-kit/node/output'
+import {AsyncLocalStorage} from 'node:async_hooks'
 
 type DevelopmentThemeId = string
 
@@ -16,6 +17,9 @@ interface ThemeStorePasswordSchema {
   [themeStore: string]: string
 }
 
+/** Preserves the theme store a command is acting on during multi environment execution */
+const themeStoreContext = new AsyncLocalStorage<{store: string}>()
+
 let _themeLocalStorageInstance: LocalStorage<ThemeLocalStorageSchema> | undefined
 let _developmentThemeLocalStorageInstance: LocalStorage<DevelopmentThemeLocalStorageSchema> | undefined
 let _replThemeLocalStorageInstance: LocalStorage<DevelopmentThemeLocalStorageSchema> | undefined
@@ -56,7 +60,8 @@ function themeStorePasswordStorage() {
 }
 
 export function getThemeStore(storage: LocalStorage<ThemeLocalStorageSchema> = themeLocalStorage()) {
-  return storage.get('themeStore')
+  const context = themeStoreContext.getStore()
+  return context?.store ? context.store : storage.get('themeStore')
 }
 
 export function setThemeStore(store: string, storage: LocalStorage<ThemeLocalStorageSchema> = themeLocalStorage()) {
@@ -144,3 +149,8 @@ function assertThemeStoreExists(storage: LocalStorage<ThemeLocalStorageSchema> =
   }
   return themeStore
 }
+
+/** Provides theme store context to each environment during concurrent multi environment execution */
+export async function useThemeStoreContext<T>(store: string, callback: () => Promise<T>): Promise<T> {
+  return themeStoreContext.run({store}, callback)
+}

packages/theme/src/cli/utilities/theme-command.ts

@@ -1,6 +1,7 @@
 import {ensureThemeStore} from './theme-store.js'
 import {configurationFileName} from '../constants.js'
 import metadata from '../metadata.js'
+import {useThemeStoreContext} from '../services/local-storage.js'
 import {hashString} from '@shopify/cli-kit/node/crypto'
 import {Input} from '@oclif/core/interfaces'
 import Command, {ArgOutput, FlagOutput} from '@shopify/cli-kit/node/base-command'
@@ -153,9 +154,8 @@ export default abstract class ThemeCommand extends Command {
     environmentMap: Map<EnvironmentName, FlagValues>,
     requiredFlags: Exclude<RequiredFlags, null>,
   ) {
-    const valid: {environment: EnvironmentName; flags: FlagValues; session: AdminSession}[] = []
+    const valid: {environment: EnvironmentName; flags: FlagValues}[] = []
     const invalid: {environment: EnvironmentName; reason: string}[] = []
-    const commandName = this.constructor.name.toLowerCase()
 
     for (const [environmentName, environmentFlags] of environmentMap) {
       const validationResult = this.validConfig(environmentFlags, requiredFlags, environmentName)
@@ -165,12 +165,7 @@ export default abstract class ThemeCommand extends Command {
         continue
       }
 
-      // eslint-disable-next-line no-await-in-loop
-      const session = await this.createSession(environmentFlags)
-
-      recordEvent(`theme-command:${commandName}:multi-env:authenticated`)
-
-      valid.push({environment: environmentName, flags: environmentFlags, session})
+      valid.push({environment: environmentName, flags: environmentFlags})
     }
 
     return {valid, invalid}
@@ -227,17 +222,23 @@ export default abstract class ThemeCommand extends Command {
    * Run the command in each valid environment concurrently
    * @param validEnvironments - The valid environments to run the command in
    */
-  private async runConcurrent(
-    validEnvironments: {environment: EnvironmentName; flags: FlagValues; session: AdminSession}[],
-  ) {
+  private async runConcurrent(validEnvironments: {environment: EnvironmentName; flags: FlagValues}[]) {
     const abortController = new AbortController()
 
     await renderConcurrent({
-      processes: validEnvironments.map(({environment, flags, session}) => ({
+      processes: validEnvironments.map(({environment, flags}) => ({
         prefix: environment,
         action: async (stdout: Writable, stderr: Writable, _signal) => {
           try {
-            await this.command(flags, session, true, {stdout, stderr})
+            const store = flags.store as string
+            await useThemeStoreContext(store, async () => {
+              const session = await this.createSession(flags)
+
+              const commandName = this.constructor.name.toLowerCase()
+              recordEvent(`theme-command:${commandName}:multi-env:authenticated`)
+
+              await this.command(flags, session, true, {stdout, stderr})
+            })
 
             // eslint-disable-next-line no-catch-all/no-catch-all
           } catch (error) {