feat(fulfillment): shipping options context field (#7169)

packages/core-flows/src/definition/cart/steps/get-shipping-option-price-sets.ts

@@ -1,79 +0,0 @@
-import { ModuleRegistrationName } from "@medusajs/modules-sdk"
-import { IPricingModuleService, PricingContext } from "@medusajs/types"
-import {
-  ContainerRegistrationKeys,
-  MedusaError,
-  arrayDifference,
-  remoteQueryObjectFromString,
-} from "@medusajs/utils"
-import { StepResponse, createStep } from "@medusajs/workflows-sdk"
-
-interface StepInput {
-  optionIds: string[]
-  context?: Record<string, unknown>
-}
-
-export const getShippingOptionPriceSetsStepId = "get-shipping-option-price-sets"
-export const getShippingOptionPriceSetsStep = createStep(
-  getShippingOptionPriceSetsStepId,
-  async (data: StepInput, { container }) => {
-    if (!data.optionIds.length) {
-      return new StepResponse({})
-    }
-
-    const pricingModuleService = container.resolve<IPricingModuleService>(
-      ModuleRegistrationName.PRICING
-    )
-
-    const remoteQuery = container.resolve(
-      ContainerRegistrationKeys.REMOTE_QUERY
-    )
-
-    const query = remoteQueryObjectFromString({
-      entryPoint: "shipping_option_price_set",
-      fields: ["id", "shipping_option_id", "price_set_id"],
-      variables: {
-        shipping_option_id: data.optionIds,
-      },
-    })
-
-    const optionPriceSets = await remoteQuery(query)
-
-    const optionsMissingPrices = arrayDifference(
-      data.optionIds,
-      optionPriceSets.map((v) => v.shipping_option_id)
-    )
-
-    if (optionsMissingPrices.length) {
-      throw new MedusaError(
-        MedusaError.Types.INVALID_DATA,
-        `Shipping options with IDs ${optionsMissingPrices.join(
-          ", "
-        )} do not have a price`
-      )
-    }
-
-    const calculatedPriceSets = await pricingModuleService.calculatePrices(
-      { id: optionPriceSets.map((v) => v.price_set_id) },
-      { context: data.context as PricingContext["context"] }
-    )
-
-    const idToPriceSet = new Map<string, Record<string, any>>(
-      calculatedPriceSets.map((p) => [p.id, p])
-    )
-
-    const optionToCalculatedPriceSets = optionPriceSets.reduce(
-      (acc, { shipping_option_id, price_set_id }) => {
-        const calculatedPriceSet = idToPriceSet.get(price_set_id)
-        if (calculatedPriceSet) {
-          acc[shipping_option_id] = calculatedPriceSet
-        }
-
-        return acc
-      },
-      {}
-    )
-
-    return new StepResponse(optionToCalculatedPriceSets)
-  }
-)

packages/core-flows/src/definition/cart/steps/index.ts

@@ -9,7 +9,6 @@ export * from "./find-or-create-customer"
 export * from "./find-sales-channel"
 export * from "./get-actions-to-compute-from-promotions"
 export * from "./get-item-tax-lines"
-export * from "./get-shipping-option-price-sets"
 export * from "./get-variant-price-sets"
 export * from "./get-variants"
 export * from "./prepare-adjustments-from-promotion-actions"

packages/core-flows/src/definition/cart/steps/link-cart-payment-collection.ts

@@ -1,4 +1,5 @@
 import { Modules } from "@medusajs/modules-sdk"
+import { ContainerRegistrationKeys } from "@medusajs/utils"
 import { StepResponse, createStep } from "@medusajs/workflows-sdk"
 
 type StepInput = {
@@ -13,7 +14,7 @@ export const linkCartAndPaymentCollectionsStepId =
 export const linkCartAndPaymentCollectionsStep = createStep(
   linkCartAndPaymentCollectionsStepId,
   async (data: StepInput, { container }) => {
-    const remoteLink = container.resolve("remoteLink")
+    const remoteLink = container.resolve(ContainerRegistrationKeys.REMOTE_LINK)
 
     const links = data.links.map((d) => ({
       [Modules.CART]: { cart_id: d.cart_id },
@@ -29,7 +30,7 @@ export const linkCartAndPaymentCollectionsStep = createStep(
       return
     }
 
-    const remoteLink = container.resolve("remoteLink")
+    const remoteLink = container.resolve(ContainerRegistrationKeys.REMOTE_LINK)
 
     const links = data.links.map((d) => ({
       [Modules.CART]: { cart_id: d.cart_id },

packages/core-flows/src/definition/cart/workflows/add-to-cart.ts

@@ -12,10 +12,7 @@ import { useRemoteQueryStep } from "../../../common/steps/use-remote-query"
 import {
   addToCartStep,
   confirmInventoryStep,
-  getVariantPriceSetsStep,
-  getVariantsStep,
   refreshCartShippingMethodsStep,
-  validateVariantsExistStep,
 } from "../steps"
 import { refreshCartPromotionsStep } from "../steps/refresh-cart-promotions"
 import { updateTaxLinesStep } from "../steps/update-tax-lines"
@@ -35,69 +32,6 @@ export const addToCartWorkflow = createWorkflow(
       return (data.input.items ?? []).map((i) => i.variant_id)
     })
 
-    validateVariantsExistStep({ variantIds })
-
-    const variants = getVariantsStep({
-      filter: { id: variantIds },
-      config: {
-        select: [
-          "id",
-          "title",
-          "sku",
-          "barcode",
-          "product.id",
-          "product.title",
-          "product.description",
-          "product.subtitle",
-          "product.thumbnail",
-          "product.type",
-          "product.collection",
-          "product.handle",
-        ],
-        relations: ["product"],
-      },
-    })
-
-    const salesChannelLocations = useRemoteQueryStep({
-      entry_point: "sales_channels",
-      fields: ["id", "name", "stock_locations.id", "stock_locations.name"],
-      variables: { id: input.cart.sales_channel_id },
-    })
-
-    const productVariantInventoryItems = useRemoteQueryStep({
-      entry_point: "product_variant_inventory_items",
-      fields: ["variant_id", "inventory_item_id", "required_quantity"],
-      variables: { variant_id: variantIds },
-    }).config({ name: "inventory-items" })
-
-    const confirmInventoryInput = transform(
-      { productVariantInventoryItems, salesChannelLocations, input, variants },
-      (data) => {
-        if (!data.salesChannelLocations.length) {
-          throw new MedusaError(
-            MedusaError.Types.INVALID_DATA,
-            `Sales channel ${data.input.cart.sales_channel_id} is not associated with any stock locations.`
-          )
-        }
-
-        const items = prepareConfirmInventoryInput({
-          product_variant_inventory_items: data.productVariantInventoryItems,
-          location_ids: data.salesChannelLocations[0].stock_locations.map(
-            (l) => l.id
-          ),
-          items: data.input.items!,
-          variants: data.variants.map((v) => ({
-            id: v.id,
-            manage_inventory: v.manage_inventory,
-          })),
-        })
-
-        return { items }
-      }
-    )
-
-    confirmInventoryStep(confirmInventoryInput)
-
     // TODO: This is on par with the context used in v1.*, but we can be more flexible.
     const pricingContext = transform({ cart: input.cart }, (data) => {
       return {
@@ -107,18 +41,115 @@ export const addToCartWorkflow = createWorkflow(
       }
     })
 
-    const priceSets = getVariantPriceSetsStep({
-      variantIds,
-      context: pricingContext,
+    const variants = useRemoteQueryStep({
+      entry_point: "variants",
+      fields: [
+        "id",
+        "title",
+        "sku",
+        "barcode",
+        "manage_inventory",
+        "product.id",
+        "product.title",
+        "product.description",
+        "product.subtitle",
+        "product.thumbnail",
+        "product.type",
+        "product.collection",
+        "product.handle",
+
+        "calculated_price.calculated_amount",
+
+        "inventory_items.inventory_item_id",
+        "inventory_items.required_quantity",
+
+        "inventory_items.inventory.location_levels.stock_locations.id",
+        "inventory_items.inventory.location_levels.stock_locations.name",
+
+        "inventory_items.inventory.location_levels.stock_locations.sales_channels.id",
+        "inventory_items.inventory.location_levels.stock_locations.sales_channels.name",
+      ],
+      variables: {
+        id: variantIds,
+        calculated_price: {
+          context: pricingContext,
+        },
+      },
+      throw_if_key_not_found: true,
     })
 
-    const lineItems = transform({ priceSets, input, variants }, (data) => {
+    const confirmInventoryInput = transform({ input, variants }, (data) => {
+      const managedVariants = data.variants.filter((v) => v.manage_inventory)
+      if (!managedVariants.length) {
+        return { items: [] }
+      }
+
+      const productVariantInventoryItems: any[] = []
+
+      const stockLocations = data.variants
+        .map((v) => v.inventory_items)
+        .flat()
+        .map((ii) => {
+          productVariantInventoryItems.push({
+            variant_id: ii.variant_id,
+            inventory_item_id: ii.inventory_item_id,
+            required_quantity: ii.required_quantity,
+          })
+
+          return ii.inventory.location_levels
+        })
+        .flat()
+        .map((ll) => ll.stock_locations)
+        .flat()
+
+      const salesChannelId = data.input.cart.sales_channel_id
+      if (salesChannelId) {
+        const salesChannels = stockLocations
+          .map((sl) => sl.sales_channels)
+          .flat()
+          .filter((sc) => sc.id === salesChannelId)
+
+        if (!salesChannels.length) {
+          throw new MedusaError(
+            MedusaError.Types.INVALID_DATA,
+            `Sales channel ${salesChannelId} is not associated with any stock locations.`
+          )
+        }
+      }
+
+      const priceNotFound: string[] = data.variants
+        .filter((v) => !v.calculated_price)
+        .map((v) => v.id)
+
+      if (priceNotFound.length) {
+        throw new MedusaError(
+          MedusaError.Types.INVALID_DATA,
+          `Variants with IDs ${priceNotFound.join(", ")} do not have a price`
+        )
+      }
+
+      const items = prepareConfirmInventoryInput({
+        product_variant_inventory_items: productVariantInventoryItems,
+        location_ids: stockLocations.map((l) => l.id),
+        items: data.input.items!,
+        variants: data.variants.map((v) => ({
+          id: v.id,
+          manage_inventory: v.manage_inventory,
+        })),
+      })
+
+      return { items }
+    })
+
+    confirmInventoryStep(confirmInventoryInput)
+
+    const lineItems = transform({ input, variants }, (data) => {
       const items = (data.input.items ?? []).map((item) => {
         const variant = data.variants.find((v) => v.id === item.variant_id)!
 
         return prepareLineItemData({
           variant: variant,
-          unitPrice: data.priceSets[item.variant_id].calculated_amount,
+          unitPrice: variant.calculated_price.calculated_amount,
           quantity: item.quantity,
           metadata: item?.metadata ?? {},
           cartId: data.input.cart.id,

packages/core-flows/src/definition/cart/workflows/create-carts.ts

@@ -14,8 +14,6 @@ import {
   findOrCreateCustomerStep,
   findSalesChannelStep,
   getVariantPriceSetsStep,
-  getVariantsStep,
-  validateVariantsExistStep,
 } from "../steps"
 import { refreshCartPromotionsStep } from "../steps/refresh-cart-promotions"
 import { updateTaxLinesStep } from "../steps/update-tax-lines"
@@ -44,77 +42,9 @@ export const createCartWorkflow = createWorkflow(
       findOrCreateCustomerStep({
         customerId: input.customer_id,
         email: input.email,
-      }),
-      validateVariantsExistStep({ variantIds })
+      })
     )
 
-    const variants = getVariantsStep({
-      filter: { id: variantIds },
-      config: {
-        select: [
-          "id",
-          "title",
-          "sku",
-          "manage_inventory",
-          "barcode",
-          "product.id",
-          "product.title",
-          "product.description",
-          "product.subtitle",
-          "product.thumbnail",
-          "product.type",
-          "product.collection",
-          "product.handle",
-        ],
-        relations: ["product"],
-      },
-    })
-
-    const salesChannelLocations = useRemoteQueryStep({
-      entry_point: "sales_channels",
-      fields: ["id", "name", "stock_locations.id", "stock_locations.name"],
-      variables: { id: salesChannel.id },
-    })
-
-    const productVariantInventoryItems = useRemoteQueryStep({
-      entry_point: "product_variant_inventory_items",
-      fields: ["variant_id", "inventory_item_id", "required_quantity"],
-      variables: { variant_id: variantIds },
-    }).config({ name: "inventory-items" })
-
-    const confirmInventoryInput = transform(
-      { productVariantInventoryItems, salesChannelLocations, input, variants },
-      (data) => {
-        // We don't want to confirm inventory if there are no items in the cart.
-        if (!data.input.items) {
-          return { items: [] }
-        }
-
-        if (!data.salesChannelLocations.length) {
-          throw new MedusaError(
-            MedusaError.Types.INVALID_DATA,
-            `Sales channel ${data.input.sales_channel_id} is not associated with any stock locations.`
-          )
-        }
-
-        const items = prepareConfirmInventoryInput({
-          product_variant_inventory_items: data.productVariantInventoryItems,
-          location_ids: data.salesChannelLocations[0].stock_locations.map(
-            (l) => l.id
-          ),
-          items: data.input.items!,
-          variants: data.variants.map((v) => ({
-            id: v.id,
-            manage_inventory: v.manage_inventory,
-          })),
-        })
-
-        return { items }
-      }
-    )
-
-    confirmInventoryStep(confirmInventoryInput)
-
     // TODO: This is on par with the context used in v1.*, but we can be more flexible.
     const pricingContext = transform(
       { input, region, customerData },
@@ -127,6 +57,111 @@ export const createCartWorkflow = createWorkflow(
       }
     )
 
+    const variants = useRemoteQueryStep({
+      entry_point: "variants",
+      fields: [
+        "id",
+        "title",
+        "sku",
+        "manage_inventory",
+        "barcode",
+        "product.id",
+        "product.title",
+        "product.description",
+        "product.subtitle",
+        "product.thumbnail",
+        "product.type",
+        "product.collection",
+        "product.handle",
+
+        "calculated_price.calculated_amount",
+
+        "inventory_items.inventory_item_id",
+        "inventory_items.required_quantity",
+
+        "inventory_items.inventory.location_levels.stock_locations.id",
+        "inventory_items.inventory.location_levels.stock_locations.name",
+
+        "inventory_items.inventory.location_levels.stock_locations.sales_channels.id",
+        "inventory_items.inventory.location_levels.stock_locations.sales_channels.name",
+      ],
+      variables: {
+        id: variantIds,
+        calculated_price: {
+          context: pricingContext,
+        },
+      },
+      throw_if_key_not_found: true,
+    })
+
+    const confirmInventoryInput = transform(
+      { input, salesChannel, variants },
+      (data) => {
+        const managedVariants = data.variants.filter((v) => v.manage_inventory)
+        if (!managedVariants.length) {
+          return { items: [] }
+        }
+
+        const productVariantInventoryItems: any[] = []
+
+        const stockLocations = managedVariants
+          .map((v) => v.inventory_items)
+          .flat()
+          .map((ii) => {
+            productVariantInventoryItems.push({
+              variant_id: ii.variant_id,
+              inventory_item_id: ii.inventory_item_id,
+              required_quantity: ii.required_quantity,
+            })
+
+            return ii.inventory.location_levels
+          })
+          .flat()
+          .map((ll) => ll.stock_locations)
+          .flat()
+
+        const salesChannelId = data.salesChannel?.id
+        if (salesChannelId) {
+          const salesChannels = stockLocations
+            .map((sl) => sl.sales_channels)
+            .flat()
+            .filter((sc) => sc.id === salesChannelId)
+
+          if (!salesChannels.length) {
+            throw new MedusaError(
+              MedusaError.Types.INVALID_DATA,
+              `Sales channel ${salesChannelId} is not associated with any stock locations.`
+            )
+          }
+        }
+
+        const priceNotFound: string[] = data.variants
+          .filter((v) => !v.calculated_price)
+          .map((v) => v.id)
+
+        if (priceNotFound.length) {
+          throw new MedusaError(
+            MedusaError.Types.INVALID_DATA,
+            `Variants with IDs ${priceNotFound.join(", ")} do not have a price`
+          )
+        }
+
+        const items = prepareConfirmInventoryInput({
+          product_variant_inventory_items: productVariantInventoryItems,
+          location_ids: stockLocations.map((l) => l.id),
+          items: data.input.items!,
+          variants: data.variants.map((v) => ({
+            id: v.id,
+            manage_inventory: v.manage_inventory,
+          })),
+        })
+
+        return { items }
+      }
+    )
+
+    confirmInventoryStep(confirmInventoryInput)
+
     const priceSets = getVariantPriceSetsStep({
       variantIds,
       context: pricingContext,

packages/core-flows/src/definition/cart/workflows/list-shipping-options-for-cart.ts

@@ -1,12 +1,11 @@
 import { ListShippingOptionsForCartWorkflowInputDTO } from "@medusajs/types"
+import { deepFlatMap, MedusaError } from "@medusajs/utils"
 import {
   createWorkflow,
   transform,
   WorkflowData,
 } from "@medusajs/workflows-sdk"
 import { useRemoteQueryStep } from "../../../common/steps/use-remote-query"
-import { listShippingOptionsForContextStep } from "../../../shipping-options"
-import { getShippingOptionPriceSetsStep } from "../steps"
 
 export const listShippingOptionsForCartWorkflowId =
   "list-shipping-options-for-cart"
@@ -15,72 +14,85 @@ export const listShippingOptionsForCartWorkflow = createWorkflow(
   (input: WorkflowData<ListShippingOptionsForCartWorkflowInputDTO>) => {
     const scLocationFulfillmentSets = useRemoteQueryStep({
       entry_point: "sales_channels",
-      fields: ["stock_locations.fulfillment_sets.id"],
-      variables: { id: input.sales_channel_id },
-    })
+      fields: [
+        "stock_locations.fulfillment_sets.id",
+        "stock_locations.fulfillment_sets.name",
+        "stock_locations.fulfillment_sets.price_type",
+        "stock_locations.fulfillment_sets.service_zone_id",
+        "stock_locations.fulfillment_sets.shipping_profile_id",
+        "stock_locations.fulfillment_sets.provider_id",
+        "stock_locations.fulfillment_sets.data",
+        "stock_locations.fulfillment_sets.amount",
 
-    const listOptionsInput = transform(
-      { scLocationFulfillmentSets, input },
-      (data) => {
-        const fulfillmentSetIds = data.scLocationFulfillmentSets
-          .map((sc) =>
-            sc.stock_locations.map((loc) =>
-              loc.fulfillment_sets.map(({ id }) => id)
-            )
-          )
-          .flat(2)
+        "stock_locations.fulfillment_sets.service_zones.shipping_options.id",
+        "stock_locations.fulfillment_sets.service_zones.shipping_options.name",
+        "stock_locations.fulfillment_sets.service_zones.shipping_options.price_type",
+        "stock_locations.fulfillment_sets.service_zones.shipping_options.service_zone_id",
+        "stock_locations.fulfillment_sets.service_zones.shipping_options.shipping_profile_id",
+        "stock_locations.fulfillment_sets.service_zones.shipping_options.provider_id",
+        "stock_locations.fulfillment_sets.service_zones.shipping_options.data",
+        "stock_locations.fulfillment_sets.service_zones.shipping_options.amount",
 
-        return {
+        "stock_locations.fulfillment_sets.service_zones.shipping_options.type.id",
+        "stock_locations.fulfillment_sets.service_zones.shipping_options.type.label",
+        "stock_locations.fulfillment_sets.service_zones.shipping_options.type.description",
+        "stock_locations.fulfillment_sets.service_zones.shipping_options.type.code",
+
+        "stock_locations.fulfillment_sets.service_zones.shipping_options.provider.id",
+        "stock_locations.fulfillment_sets.service_zones.shipping_options.provider.is_enabled",
+
+        "stock_locations.fulfillment_sets.service_zones.shipping_options.calculated_price.calculated_amount",
+      ],
+      variables: {
+        id: input.sales_channel_id,
+        "stock_locations.fulfillment_sets.service_zones.shipping_options": {
           context: {
-            fulfillment_set_id: fulfillmentSetIds,
             address: {
-              city: data.input.shipping_address?.city,
-              country_code: data.input.shipping_address?.country_code,
-              province_code: data.input.shipping_address?.province,
+              city: input.shipping_address?.city,
+              country_code: input.shipping_address?.country_code,
+              province_code: input.shipping_address?.province,
             },
           },
-          config: {
-            select: [
-              "id",
-              "name",
-              "price_type",
-              "service_zone_id",
-              "shipping_profile_id",
-              "provider_id",
-              "data",
-              "amount",
-            ],
-            relations: ["type", "provider"],
+        },
+        "stock_locations.fulfillment_sets.service_zones.shipping_options.calculated_price":
+          {
+            context: {
+              currency_code: input.currency_code,
+            },
           },
-        }
-      }
-    )
-
-    const options = listShippingOptionsForContextStep(listOptionsInput)
-
-    const optionIds = transform({ options }, (data) =>
-      data.options.map((option) => option.id)
-    )
-
-    // TODO: Separate shipping options based on price_type, flat_rate vs calculated
-    const priceSets = getShippingOptionPriceSetsStep({
-      optionIds,
-      context: {
-        currency_code: input.currency_code,
       },
     })
 
     const shippingOptionsWithPrice = transform(
-      { priceSets, options },
+      { options: scLocationFulfillmentSets },
       (data) => {
-        const options = data.options.map((option) => {
-          const price = data.priceSets?.[option.id].calculated_amount
+        const optionsMissingPrices: string[] = []
 
-          return {
-            ...option,
-            amount: price,
+        const options = deepFlatMap(
+          data.options,
+          "stock_locations.fulfillment_sets.service_zones.shipping_options.calculated_price",
+          ({ shipping_options }) => {
+            const { calculated_price, ...options } = shipping_options ?? {}
+
+            if (!calculated_price) {
+              optionsMissingPrices.push(options.id)
+            }
+
+            return {
+              ...options,
+              amount: calculated_price?.calculated_amount,
+            }
           }
-        })
+        )
+
+        if (optionsMissingPrices.length) {
+          throw new MedusaError(
+            MedusaError.Types.INVALID_DATA,
+            `Shipping options with IDs ${optionsMissingPrices.join(
+              ", "
+            )} do not have a price`
+          )
+        }
 
         return options
       }

packages/core-flows/src/definition/cart/workflows/refresh-payment-collection.ts

@@ -3,6 +3,7 @@ import {
   WorkflowData,
   createStep,
   createWorkflow,
+  transform,
 } from "@medusajs/workflows-sdk"
 import { useRemoteQueryStep } from "../../../common/steps/use-remote-query"
 import {
@@ -49,21 +50,20 @@ export const refreshPaymentCollectionForCartWorkflow = createWorkflow(
         "payment_collection.payment_sessions.id",
       ],
       variables: { id: input.cart_id },
+      throw_if_key_not_found: true,
     })
 
+    const cart = transform({ carts }, (data) => data.carts[0])
+
     deletePaymentSessionStep({
-      payment_session_id: carts[0].payment_collection.payment_sessions?.[0].id,
+      payment_session_id: cart.payment_collection.payment_sessions?.[0].id,
     })
 
-    // TODO: Temporary fixed cart total, so we can test the workflow.
-    //  This will be removed when the totals utilities are built.
-    const cartTotal = 4242
-
     updatePaymentCollectionStep({
-      selector: { id: carts[0].payment_collection.id },
+      selector: { id: cart.payment_collection.id },
       update: {
-        amount: cartTotal,
-        currency_code: carts[0].currency_code,
+        amount: cart.total,
+        currency_code: cart.currency_code,
       },
     })
   }

packages/core-flows/src/definition/cart/workflows/update-line-item-in-cart.ts

@@ -7,12 +7,7 @@ import {
 import { MedusaError } from "medusa-core-utils"
 import { useRemoteQueryStep } from "../../../common/steps/use-remote-query"
 import { updateLineItemsStep } from "../../line-item/steps"
-import {
-  confirmInventoryStep,
-  getVariantPriceSetsStep,
-  getVariantsStep,
-  refreshCartShippingMethodsStep,
-} from "../steps"
+import { confirmInventoryStep, refreshCartShippingMethodsStep } from "../steps"
 import { refreshCartPromotionsStep } from "../steps/refresh-cart-promotions"
 import { cartFieldsForRefreshSteps } from "../utils/fields"
 import { prepareConfirmInventoryInput } from "../utils/prepare-confirm-inventory-input"
@@ -25,9 +20,12 @@ export const updateLineItemInCartWorkflowId = "update-line-item-in-cart"
 export const updateLineItemInCartWorkflow = createWorkflow(
   updateLineItemInCartWorkflowId,
   (input: WorkflowData<UpdateLineItemInCartWorkflowInputDTO>) => {
-    const item = transform({ input }, (data) => data.input.item)
+    const variantIds = transform({ input }, (data) => {
+      return [data.input.item.variant_id]
+    })
 
-    const pricingContext = transform({ cart: input.cart, item }, (data) => {
+    // TODO: This is on par with the context used in v1.*, but we can be more flexible.
+    const pricingContext = transform({ cart: input.cart }, (data) => {
       return {
         currency_code: data.cart.currency_code,
         region_id: data.cart.region_id,
@@ -35,78 +33,124 @@ export const updateLineItemInCartWorkflow = createWorkflow(
       }
     })
 
-    const variantIds = transform({ input }, (data) => [
-      data.input.item.variant_id!,
-    ])
+    const variants = useRemoteQueryStep({
+      entry_point: "variants",
+      fields: [
+        "id",
+        "title",
+        "sku",
+        "barcode",
+        "manage_inventory",
+        "product.id",
+        "product.title",
+        "product.description",
+        "product.subtitle",
+        "product.thumbnail",
+        "product.type",
+        "product.collection",
+        "product.handle",
 
-    const salesChannelLocations = useRemoteQueryStep({
-      entry_point: "sales_channels",
-      fields: ["id", "name", "stock_locations.id", "stock_locations.name"],
-      variables: { id: input.cart.sales_channel_id },
+        "calculated_price.calculated_amount",
+
+        "inventory_items.inventory_item_id",
+        "inventory_items.required_quantity",
+
+        "inventory_items.inventory.location_levels.stock_locations.id",
+        "inventory_items.inventory.location_levels.stock_locations.name",
+
+        "inventory_items.inventory.location_levels.stock_locations.sales_channels.id",
+        "inventory_items.inventory.location_levels.stock_locations.sales_channels.name",
+      ],
+      variables: {
+        id: variantIds,
+        calculated_price: {
+          context: pricingContext,
+        },
+      },
+      throw_if_key_not_found: true,
     })
 
-    const productVariantInventoryItems = useRemoteQueryStep({
-      entry_point: "product_variant_inventory_items",
-      fields: ["variant_id", "inventory_item_id", "required_quantity"],
-      variables: { variant_id: variantIds },
-    }).config({ name: "inventory-items" })
+    const confirmInventoryInput = transform({ input, variants }, (data) => {
+      const managedVariants = data.variants.filter((v) => v.manage_inventory)
+      if (!managedVariants.length) {
+        return { items: [] }
+      }
 
-    const variants = getVariantsStep({
-      filter: { id: variantIds },
-      config: { select: ["id", "manage_inventory"] },
-    })
+      const productVariantInventoryItems: any[] = []
 
-    const confirmInventoryInput = transform(
-      { productVariantInventoryItems, salesChannelLocations, input, variants },
-      (data) => {
-        if (!data.salesChannelLocations.length) {
+      const stockLocations = data.variants
+        .map((v) => v.inventory_items)
+        .flat()
+        .map((ii) => {
+          productVariantInventoryItems.push({
+            variant_id: ii.variant_id,
+            inventory_item_id: ii.inventory_item_id,
+            required_quantity: ii.required_quantity,
+          })
+
+          return ii.inventory.location_levels
+        })
+        .flat()
+        .map((ll) => ll.stock_locations)
+        .flat()
+
+      const salesChannelId = data.input.cart.sales_channel_id
+      if (salesChannelId) {
+        const salesChannels = stockLocations
+          .map((sl) => sl.sales_channels)
+          .flat()
+          .filter((sc) => sc.id === salesChannelId)
+
+        if (!salesChannels.length) {
           throw new MedusaError(
             MedusaError.Types.INVALID_DATA,
-            `Sales channel ${data.input.cart.sales_channel_id} is not associated with any stock locations.`
+            `Sales channel ${salesChannelId} is not associated with any stock locations.`
           )
         }
-
-        const items = prepareConfirmInventoryInput({
-          product_variant_inventory_items: data.productVariantInventoryItems,
-          location_ids: data.salesChannelLocations[0].stock_locations.map(
-            (l) => l.id
-          ),
-          items: [data.input.item],
-          variants: data.variants.map((v) => ({
-            id: v.id,
-            manage_inventory: v.manage_inventory,
-          })),
-        })
-
-        return { items }
       }
-    )
+
+      const priceNotFound: string[] = data.variants
+        .filter((v) => !v.calculated_price)
+        .map((v) => v.id)
+
+      if (priceNotFound.length) {
+        throw new MedusaError(
+          MedusaError.Types.INVALID_DATA,
+          `Variants with IDs ${priceNotFound.join(", ")} do not have a price`
+        )
+      }
+
+      const items = prepareConfirmInventoryInput({
+        product_variant_inventory_items: productVariantInventoryItems,
+        location_ids: stockLocations.map((l) => l.id),
+        items: [data.input.item!],
+        variants: data.variants.map((v) => ({
+          id: v.id,
+          manage_inventory: v.manage_inventory,
+        })),
+      })
+
+      return { items }
+    })
 
     confirmInventoryStep(confirmInventoryInput)
 
-    const priceSets = getVariantPriceSetsStep({
-      variantIds,
-      context: pricingContext,
-    })
-
-    const lineItemUpdate = transform({ input, priceSets, item }, (data) => {
-      const price = data.priceSets[data.item.variant_id!].calculated_amount
+    const lineItemUpdate = transform({ input, variants }, (data) => {
+      const variant = data.variants[0]
+      const item = data.input.item
 
       return {
         data: {
           ...data.input.update,
-          unit_price: price,
+          unit_price: variant.calculated_price.calculated_amount,
         },
         selector: {
-          id: data.input.item.id,
+          id: item.id,
         },
       }
     })
 
-    const result = updateLineItemsStep({
-      data: lineItemUpdate.data,
-      selector: lineItemUpdate.selector,
-    })
+    const result = updateLineItemsStep(lineItemUpdate)
 
     const cart = useRemoteQueryStep({
       entry_point: "cart",