feat: Set pricing context hook (#11891)

.changeset/flat-meals-rhyme.md

@@ -0,0 +1,7 @@
+---
+"@medusajs/core-flows": patch
+"@medusajs/workflows-sdk": patch
+"integration-tests-modules": patch
+---
+
+feat: Set pricing context hook

packages/core/core-flows/src/cart/utils/schemas.ts

@@ -0,0 +1,2 @@
+import z from "zod"
+export const pricingContextResult = z.record(z.string(), z.any()).optional()

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

@@ -1,4 +1,5 @@
 import {
+  AdditionalData,
   AddToCartWorkflowInputDTO,
   ConfirmVariantInventoryWorkflowInputDTO,
 } from "@medusajs/framework/types"
@@ -33,6 +34,7 @@ import {
 } from "../utils/prepare-line-item-data"
 import { confirmVariantInventoryWorkflow } from "./confirm-variant-inventory"
 import { refreshCartItemsWorkflow } from "./refresh-cart-items"
+import { pricingContextResult } from "../utils/schemas"
 
 const cartFields = ["completed_at"].concat(cartFieldsForPricingContext)
 
@@ -71,7 +73,7 @@ export const addToCartWorkflowId = "add-to-cart"
  */
 export const addToCartWorkflow = createWorkflow(
   addToCartWorkflowId,
-  (input: WorkflowData<AddToCartWorkflowInputDTO>) => {
+  (input: WorkflowData<AddToCartWorkflowInputDTO & AdditionalData>) => {
     const cartQuery = useQueryGraphStep({
       entity: "cart",
       filters: { id: input.cart_id },
@@ -93,6 +95,35 @@ export const addToCartWorkflow = createWorkflow(
       return (data.input.items ?? []).map((i) => i.variant_id).filter(Boolean)
     })
 
+    const setPricingContext = createHook(
+      "setPricingContext",
+      {
+        cart,
+        variantIds,
+        items: input.items,
+        additional_data: input.additional_data,
+      },
+      {
+        resultValidator: pricingContextResult,
+      }
+    )
+
+    const setPricingContextResult = setPricingContext.getResult()
+    const pricingContext = transform(
+      { cart, setPricingContextResult },
+      (data) => {
+        return {
+          ...data.cart,
+          ...(data.setPricingContextResult ? data.setPricingContextResult : {}),
+          currency_code: data.cart.currency_code,
+          region_id: data.cart.region_id,
+          region: data.cart.region,
+          customer_id: data.cart.customer_id,
+          customer: data.cart.customer,
+        }
+      }
+    )
+
     const variants = when({ variantIds }, ({ variantIds }) => {
       return !!variantIds.length
     }).then(() => {
@@ -102,7 +133,7 @@ export const addToCartWorkflow = createWorkflow(
         variables: {
           id: variantIds,
           calculated_price: {
-            context: cart,
+            context: pricingContext,
           },
         },
       })
@@ -201,7 +232,7 @@ export const addToCartWorkflow = createWorkflow(
     })
 
     return new WorkflowResponse(void 0, {
-      hooks: [validate],
+      hooks: [validate, setPricingContext] as const,
     })
   }
 )

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

@@ -36,6 +36,7 @@ import { refreshPaymentCollectionForCartWorkflow } from "./refresh-payment-colle
 import { updateCartPromotionsWorkflow } from "./update-cart-promotions"
 import { updateTaxLinesWorkflow } from "./update-tax-lines"
 import { validateSalesChannelStep } from "../steps/validate-sales-channel"
+import { pricingContextResult } from "../utils/schemas"
 
 /**
  * The data to create the cart, along with custom data that's passed to the workflow's hooks.
@@ -98,16 +99,31 @@ export const createCartWorkflow = createWorkflow(
     )
 
     validateSalesChannelStep({ salesChannel })
+    const setPricingContext = createHook(
+      "setPricingContext",
+      {
+        region,
+        variantIds,
+        salesChannel,
+        customerData,
+        additional_data: input.additional_data,
+      },
+      {
+        resultValidator: pricingContextResult,
+      }
+    )
+    const setPricingContextResult = setPricingContext.getResult()
 
     // TODO: This is on par with the context used in v1.*, but we can be more flexible.
     const pricingContext = transform(
-      { input, region, customerData },
+      { input, region, customerData, setPricingContextResult },
       (data) => {
         if (!data.region) {
           throw new MedusaError(MedusaError.Types.NOT_FOUND, "No regions found")
         }
 
         return {
+          ...(data.setPricingContextResult ? data.setPricingContextResult : {}),
           currency_code: data.input.currency_code ?? data.region.currency_code,
           region_id: data.region.id,
           customer_id: data.customerData.customer?.id,
@@ -248,7 +264,7 @@ export const createCartWorkflow = createWorkflow(
     })
 
     return new WorkflowResponse(cart, {
-      hooks: [validate, cartCreated],
+      hooks: [validate, cartCreated, setPricingContext] as const,
     })
   }
 )

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

@@ -1,4 +1,5 @@
 import {
+  createHook,
   createWorkflow,
   transform,
   WorkflowData,
@@ -7,8 +8,12 @@ import {
 import { useQueryGraphStep, validatePresenceOfStep } from "../../common"
 import { useRemoteQueryStep } from "../../common/steps/use-remote-query"
 import { cartFieldsForPricingContext } from "../utils/fields"
-import { ListShippingOptionsForCartWorkflowInput } from "@medusajs/types"
+import {
+  AdditionalData,
+  ListShippingOptionsForCartWorkflowInput,
+} from "@medusajs/types"
 import { isDefined } from "@medusajs/framework/utils"
+import { pricingContextResult } from "../utils/schemas"
 
 export const listShippingOptionsForCartWorkflowId =
   "list-shipping-options-for-cart"
@@ -41,7 +46,11 @@ export const listShippingOptionsForCartWorkflowId =
  */
 export const listShippingOptionsForCartWorkflow = createWorkflow(
   listShippingOptionsForCartWorkflowId,
-  (input: WorkflowData<ListShippingOptionsForCartWorkflowInput>) => {
+  (
+    input: WorkflowData<
+      ListShippingOptionsForCartWorkflowInput & AdditionalData
+    >
+  ) => {
     const cartQuery = useQueryGraphStep({
       entity: "cart",
       filters: { id: input.cart_id },
@@ -96,9 +105,22 @@ export const listShippingOptionsForCartWorkflow = createWorkflow(
       }
     )
 
+    const setPricingContext = createHook(
+      "setPricingContext",
+      {
+        cart: cart,
+        fulfillmentSetIds,
+        additional_data: input.additional_data,
+      },
+      {
+        resultValidator: pricingContextResult,
+      }
+    )
+    const setPricingContextResult = setPricingContext.getResult()
+
     const queryVariables = transform(
-      { input, fulfillmentSetIds, cart },
+      { input, fulfillmentSetIds, cart, setPricingContextResult },
-      ({ input, fulfillmentSetIds, cart }) => {
+      ({ input, fulfillmentSetIds, cart, setPricingContextResult }) => {
         return {
           id: input.option_ids,
 
@@ -122,7 +144,17 @@ export const listShippingOptionsForCartWorkflow = createWorkflow(
             },
           },
 
-          calculated_price: { context: cart },
+          calculated_price: {
+            context: {
+              ...cart,
+              ...(setPricingContextResult ? setPricingContextResult : {}),
+              currency_code: cart.currency_code,
+              region_id: cart.region_id,
+              region: cart.region,
+              customer_id: cart.customer_id,
+              customer: cart.customer,
+            },
+          },
         }
       }
     )
@@ -202,6 +234,8 @@ export const listShippingOptionsForCartWorkflow = createWorkflow(
         })
     )
 
-    return new WorkflowResponse(shippingOptionsWithPrice)
+    return new WorkflowResponse(shippingOptionsWithPrice, {
+      hooks: [setPricingContext] as const,
+    })
   }
 )

packages/core/core-flows/src/cart/workflows/refresh-cart-items.ts

@@ -28,6 +28,8 @@ import { refreshPaymentCollectionForCartWorkflow } from "./refresh-payment-colle
 import { updateCartPromotionsWorkflow } from "./update-cart-promotions"
 import { updateTaxLinesWorkflow } from "./update-tax-lines"
 import { upsertTaxLinesWorkflow } from "./upsert-tax-lines"
+import { AdditionalData } from "@medusajs/types"
+import { pricingContextResult } from "../utils/schemas"
 
 /**
  * The details of the cart to refresh.
@@ -91,7 +93,20 @@ export const refreshCartItemsWorkflowId = "refresh-cart-items"
  */
 export const refreshCartItemsWorkflow = createWorkflow(
   refreshCartItemsWorkflowId,
-  (input: WorkflowData<RefreshCartItemsWorkflowInput>) => {
+  (input: WorkflowData<RefreshCartItemsWorkflowInput & AdditionalData>) => {
+    const setPricingContext = createHook(
+      "setPricingContext",
+      {
+        cart_id: input.cart_id,
+        items: input.items,
+        additional_data: input.additional_data,
+      },
+      {
+        resultValidator: pricingContextResult,
+      }
+    )
+    const setPricingContextResult = setPricingContext.getResult()
+
     when({ input }, ({ input }) => {
       return !!input.force_refresh
     }).then(() => {
@@ -106,9 +121,22 @@ export const refreshCartItemsWorkflow = createWorkflow(
         return (data.cart.items ?? []).map((i) => i.variant_id).filter(Boolean)
       })
 
-      const cartPricingContext = transform({ cart }, ({ cart }) => {
+      const cartPricingContext = transform(
-        return filterObjectByKeys(cart, cartFieldsForPricingContext)
+        { cart, setPricingContextResult },
-      })
+        (data) => {
+          return {
+            ...filterObjectByKeys(data.cart, cartFieldsForPricingContext),
+            ...(data.setPricingContextResult
+              ? data.setPricingContextResult
+              : {}),
+            currency_code: data.cart.currency_code,
+            region_id: data.cart.region_id,
+            region: data.cart.region,
+            customer_id: data.cart.customer_id,
+            customer: data.cart.customer,
+          }
+        }
+      )
 
       const variants = useRemoteQueryStep({
         entry_point: "variants",
@@ -235,6 +263,8 @@ export const refreshCartItemsWorkflow = createWorkflow(
       input: { cart_id: input.cart_id },
     })
 
-    return new WorkflowResponse(refetchedCart)
+    return new WorkflowResponse(refetchedCart, {
+      hooks: [setPricingContext] as const,
+    })
   }
 )

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

@@ -1,4 +1,7 @@
-import { UpdateLineItemInCartWorkflowInputDTO } from "@medusajs/framework/types"
+import {
+  AdditionalData,
+  UpdateLineItemInCartWorkflowInputDTO,
+} from "@medusajs/framework/types"
 import { isDefined, MedusaError } from "@medusajs/framework/utils"
 import {
   createHook,
@@ -19,6 +22,7 @@ import {
 } from "../utils/fields"
 import { confirmVariantInventoryWorkflow } from "./confirm-variant-inventory"
 import { refreshCartItemsWorkflow } from "./refresh-cart-items"
+import { pricingContextResult } from "../utils/schemas"
 
 const cartFields = cartFieldsForPricingContext.concat(["items.*"])
 
@@ -49,7 +53,9 @@ export const updateLineItemInCartWorkflowId = "update-line-item-in-cart"
  */
 export const updateLineItemInCartWorkflow = createWorkflow(
   updateLineItemInCartWorkflowId,
-  (input: WorkflowData<UpdateLineItemInCartWorkflowInputDTO>) => {
+  (
+    input: WorkflowData<UpdateLineItemInCartWorkflowInputDTO & AdditionalData>
+  ) => {
     const cartQuery = useQueryGraphStep({
       entity: "cart",
       filters: { id: input.cart_id },
@@ -73,6 +79,35 @@ export const updateLineItemInCartWorkflow = createWorkflow(
       return [item.variant_id].filter(Boolean)
     })
 
+    const setPricingContext = createHook(
+      "setPricingContext",
+      {
+        cart,
+        item,
+        variantIds,
+        additional_data: input.additional_data,
+      },
+      {
+        resultValidator: pricingContextResult,
+      }
+    )
+
+    const setPricingContextResult = setPricingContext.getResult()
+    const pricingContext = transform(
+      { cart, setPricingContextResult },
+      (data) => {
+        return {
+          ...data.cart,
+          ...(data.setPricingContextResult ? data.setPricingContextResult : {}),
+          currency_code: data.cart.currency_code,
+          region_id: data.cart.region_id,
+          region: data.cart.region,
+          customer_id: data.cart.customer_id,
+          customer: data.cart.customer,
+        }
+      }
+    )
+
     const variants = when({ variantIds }, ({ variantIds }) => {
       return !!variantIds.length
     }).then(() => {
@@ -82,7 +117,7 @@ export const updateLineItemInCartWorkflow = createWorkflow(
         variables: {
           id: variantIds,
           calculated_price: {
-            context: cart,
+            context: pricingContext,
           },
         },
       }).config({ name: "fetch-variants" })
@@ -98,7 +133,7 @@ export const updateLineItemInCartWorkflow = createWorkflow(
 
     confirmVariantInventoryWorkflow.runAsStep({
       input: {
-        sales_channel_id: cart.sales_channel_id,
+        sales_channel_id: pricingContext.sales_channel_id,
         variants,
         items,
       },
@@ -147,7 +182,7 @@ export const updateLineItemInCartWorkflow = createWorkflow(
     })
 
     return new WorkflowResponse(void 0, {
-      hooks: [validate],
+      hooks: [validate, setPricingContext] as const,
     })
   }
 )

packages/core/core-flows/src/order/utils/fetch-shipping-option.ts

@@ -1,4 +1,5 @@
 import {
+  AdditionalData,
   BigNumberInput,
   CalculatedRMAShippingContext,
   CalculateShippingOptionPriceDTO,
@@ -6,6 +7,7 @@ import {
 } from "@medusajs/framework/types"
 import {
   WorkflowResponse,
+  createHook,
   createWorkflow,
   transform,
   when,
@@ -13,6 +15,7 @@ import {
 import { BigNumber, ShippingOptionPriceType } from "@medusajs/framework/utils"
 import { calculateShippingOptionsPricesStep } from "../../fulfillment/steps"
 import { useRemoteQueryStep } from "../../common"
+import { pricingContextResult } from "../../cart/utils/schemas"
 
 const COMMON_OPTIONS_FIELDS = [
   "id",
@@ -90,9 +93,7 @@ export const createOrderEditShippingMethodWorkflowId = "fetch-shipping-option"
  */
 export const fetchShippingOptionForOrderWorkflow = createWorkflow(
   createOrderEditShippingMethodWorkflowId,
-  function (
+  function (input: FetchShippingOptionForOrderWorkflowInput & AdditionalData) {
-    input: FetchShippingOptionForOrderWorkflowInput
-  ): WorkflowResponse<FetchShippingOptionForOrderWorkflowOutput> {
     const initialOption = useRemoteQueryStep({
       entry_point: "shipping_option",
       variables: { id: input.shipping_option_id },
@@ -170,6 +171,20 @@ export const fetchShippingOptionForOrderWorkflow = createWorkflow(
       return shippingOptionsWithPrice
     })
 
+    const setPricingContext = createHook("setPricingContext", input, {
+      resultValidator: pricingContextResult,
+    })
+    const setPricingContextResult = setPricingContext.getResult()
+    const pricingContext = transform(
+      { input, setPricingContextResult },
+      (data) => {
+        return {
+          ...(data.setPricingContextResult ? data.setPricingContextResult : {}),
+          currency_code: data.input.currency_code,
+        }
+      }
+    )
+
     const flatRateShippingOption = when(
       "option-flat",
       { isCalculatedPriceShippingOption },
@@ -186,7 +201,7 @@ export const fetchShippingOptionForOrderWorkflow = createWorkflow(
         variables: {
           id: input.shipping_option_id,
           calculated_price: {
-            context: { currency_code: input.currency_code },
+            context: pricingContext,
           },
         },
         list: false,
@@ -195,7 +210,7 @@ export const fetchShippingOptionForOrderWorkflow = createWorkflow(
       return shippingOption
     })
 
-    const result = transform(
+    const result: FetchShippingOptionForOrderWorkflowOutput = transform(
       {
         calculatedPriceShippingOption,
         flatRateShippingOption,
@@ -205,6 +220,6 @@ export const fetchShippingOptionForOrderWorkflow = createWorkflow(
       }
     )
 
-    return new WorkflowResponse(result)
+    return new WorkflowResponse(result, { hooks: [setPricingContext] as const })
   }
 )

packages/core/core-flows/src/order/workflows/add-line-items.ts

@@ -1,6 +1,11 @@
-import { OrderLineItemDTO, OrderWorkflow } from "@medusajs/framework/types"
+import {
+  AdditionalData,
+  OrderLineItemDTO,
+  OrderWorkflow,
+} from "@medusajs/framework/types"
 import { isDefined, MedusaError } from "@medusajs/framework/utils"
 import {
+  createHook,
   createWorkflow,
   parallelize,
   transform,
@@ -21,6 +26,7 @@ import { confirmVariantInventoryWorkflow } from "../../cart/workflows/confirm-va
 import { useRemoteQueryStep } from "../../common"
 import { createOrderLineItemsStep } from "../steps"
 import { productVariantsFields } from "../utils/fields"
+import { pricingContextResult } from "../../cart/utils/schemas"
 
 function prepareLineItems(data) {
   const items = (data.input.items ?? []).map((item) => {
@@ -82,8 +88,10 @@ export const addOrderLineItemsWorkflowId = "order-add-line-items"
 export const addOrderLineItemsWorkflow = createWorkflow(
   addOrderLineItemsWorkflowId,
   (
-    input: WorkflowData<OrderWorkflow.OrderAddLineItemWorkflowInput>
+    input: WorkflowData<
-  ): WorkflowResponse<OrderAddLineItemWorkflowOutput> => {
+      OrderWorkflow.OrderAddLineItemWorkflowInput & AdditionalData
+    >
+  ) => {
     const order = useRemoteQueryStep({
       entry_point: "orders",
       fields: [
@@ -118,14 +126,30 @@ export const addOrderLineItemsWorkflow = createWorkflow(
       })
     )
 
+    const setPricingContext = createHook(
+      "setPricingContext",
+      {
+        order,
+        variantIds,
+        region,
+        customerData,
+        additional_data: input.additional_data,
+      },
+      {
+        resultValidator: pricingContextResult,
+      }
+    )
+    const setPricingContextResult = setPricingContext.getResult()
+
     const pricingContext = transform(
-      { input, region, customerData, order },
+      { input, region, customerData, order, setPricingContextResult },
       (data) => {
         if (!data.region) {
           throw new MedusaError(MedusaError.Types.NOT_FOUND, "Region not found")
         }
 
         return {
+          ...(data.setPricingContextResult ? data.setPricingContextResult : {}),
           currency_code: data.order.currency_code ?? data.region.currency_code,
           region_id: data.region.id,
           customer_id: data.customerData.customer?.id,
@@ -165,7 +189,10 @@ export const addOrderLineItemsWorkflow = createWorkflow(
     return new WorkflowResponse(
       createOrderLineItemsStep({
         items: lineItems,
-      })
+      }) satisfies OrderAddLineItemWorkflowOutput,
+      {
+        hooks: [setPricingContext] as const,
+      }
     )
   }
 )

packages/core/core-flows/src/order/workflows/claim/update-claim-shipping-method.ts

@@ -1,4 +1,5 @@
 import {
+  AdditionalData,
   OrderChangeActionDTO,
   OrderChangeDTO,
   OrderClaimDTO,
@@ -9,6 +10,7 @@ import { ChangeActionType, OrderChangeStatus } from "@medusajs/framework/utils"
 import {
   WorkflowData,
   WorkflowResponse,
+  createHook,
   createStep,
   createWorkflow,
   parallelize,
@@ -26,6 +28,7 @@ import {
   throwIfOrderChangeIsNotActive,
 } from "../../utils/order-validation"
 import { prepareShippingMethodUpdate } from "../../utils/prepare-shipping-method"
+import { pricingContextResult } from "../../../cart/utils/schemas"
 
 /**
  * The data to validate that a claim's shipping method can be updated.
@@ -42,7 +45,10 @@ export type UpdateClaimShippingMethodValidationStepInput = {
   /**
    * The details of updating the shipping method.
    */
-  input: Pick<OrderWorkflow.UpdateClaimShippingMethodWorkflowInput, "claim_id" | "action_id">
+  input: Pick<
+    OrderWorkflow.UpdateClaimShippingMethodWorkflowInput,
+    "claim_id" | "action_id"
+  >
 }
 
 /**
@@ -128,8 +134,10 @@ export const updateClaimShippingMethodWorkflowId =
 export const updateClaimShippingMethodWorkflow = createWorkflow(
   updateClaimShippingMethodWorkflowId,
   function (
-    input: WorkflowData<OrderWorkflow.UpdateClaimShippingMethodWorkflowInput>
+    input: WorkflowData<
-  ): WorkflowResponse<OrderPreviewDTO> {
+      OrderWorkflow.UpdateClaimShippingMethodWorkflowInput & AdditionalData
+    >
+  ) {
     const orderClaim: OrderClaimDTO = useRemoteQueryStep({
       entry_point: "order_claim",
       fields: [
@@ -157,6 +165,19 @@ export const updateClaimShippingMethodWorkflow = createWorkflow(
       list: false,
     }).config({ name: "order-change-query" })
 
+    const setPricingContext = createHook(
+      "setPricingContext",
+      {
+        order_claim: orderClaim,
+        order_change: orderChange,
+        additional_data: input.additional_data,
+      },
+      {
+        resultValidator: pricingContextResult,
+      }
+    )
+    const setPricingContextResult = setPricingContext.getResult()
+
     const shippingOptions = when({ input }, ({ input }) => {
       return input.data?.custom_amount === null
     }).then(() => {
@@ -174,6 +195,18 @@ export const updateClaimShippingMethodWorkflow = createWorkflow(
         }
       )
 
+      const pricingContext = transform(
+        { action, setPricingContextResult },
+        (data) => {
+          return {
+            ...(data.setPricingContextResult
+              ? data.setPricingContextResult
+              : {}),
+            currency_code: data.action.currency_code,
+          }
+        }
+      )
+
       const shippingMethod = useRemoteQueryStep({
         entry_point: "order_shipping_method",
         fields: ["id", "shipping_option_id"],
@@ -194,7 +227,7 @@ export const updateClaimShippingMethodWorkflow = createWorkflow(
         variables: {
           id: shippingMethod.shipping_option_id,
           calculated_price: {
-            context: { currency_code: action.currency_code },
+            context: pricingContext,
           },
         },
       }).config({ name: "fetch-shipping-option" })
@@ -212,6 +245,11 @@ export const updateClaimShippingMethodWorkflow = createWorkflow(
       updateOrderShippingMethodsStep([updateData.shippingMethod!])
     )
 
-    return new WorkflowResponse(previewOrderChangeStep(orderClaim.order_id))
+    return new WorkflowResponse(
+      previewOrderChangeStep(orderClaim.order_id) as OrderPreviewDTO,
+      {
+        hooks: [setPricingContext] as const,
+      }
+    )
   }
 )

packages/core/core-flows/src/order/workflows/create-order.ts

@@ -23,6 +23,7 @@ import { useRemoteQueryStep } from "../../common"
 import { createOrdersStep } from "../steps"
 import { productVariantsFields } from "../utils/fields"
 import { updateOrderTaxLinesWorkflow } from "./update-tax-lines"
+import { pricingContextResult } from "../../cart/utils/schemas"
 
 function prepareLineItems(data) {
   const items = (data.input.items ?? []).map((item) => {
@@ -150,15 +151,30 @@ export const createOrderWorkflow = createWorkflow(
       })
     )
 
+    const setPricingContext = createHook(
+      "setPricingContext",
+      {
+        variantIds,
+        region,
+        customerData,
+        additional_data: input.additional_data,
+      },
+      {
+        resultValidator: pricingContextResult,
+      }
+    )
+    const setPricingContextResult = setPricingContext.getResult()
+
     // TODO: This is on par with the context used in v1.*, but we can be more flexible.
     const pricingContext = transform(
-      { input, region, customerData },
+      { input, region, customerData, setPricingContextResult },
       (data) => {
         if (!data.region) {
           throw new MedusaError(MedusaError.Types.NOT_FOUND, "Region not found")
         }
 
         return {
+          ...(data.setPricingContextResult ? data.setPricingContextResult : {}),
           currency_code: data.input.currency_code ?? data.region.currency_code,
           region_id: data.region.id,
           customer_id: data.customerData.customer?.id,
@@ -222,7 +238,7 @@ export const createOrderWorkflow = createWorkflow(
     })
 
     return new WorkflowResponse(order, {
-      hooks: [orderCreated],
+      hooks: [orderCreated, setPricingContext] as const,
     })
   }
 )

packages/core/core-flows/src/order/workflows/exchange/update-exchange-shipping-method.ts

@@ -1,4 +1,5 @@
 import {
+  AdditionalData,
   OrderChangeActionDTO,
   OrderChangeDTO,
   OrderExchangeDTO,
@@ -9,6 +10,7 @@ import { ChangeActionType, OrderChangeStatus } from "@medusajs/framework/utils"
 import {
   WorkflowData,
   WorkflowResponse,
+  createHook,
   createStep,
   createWorkflow,
   parallelize,
@@ -26,6 +28,7 @@ import {
   throwIfOrderChangeIsNotActive,
 } from "../../utils/order-validation"
 import { prepareShippingMethodUpdate } from "../../utils/prepare-shipping-method"
+import { pricingContextResult } from "../../../cart/utils/schemas"
 
 /**
  * The data to validate that an exchange's shipping method can be updated.
@@ -42,7 +45,10 @@ export type UpdateExchangeShippingMethodValidationStepInput = {
   /**
    * The details of the shipping method update.
    */
-  input: Pick<OrderWorkflow.UpdateExchangeShippingMethodWorkflowInput, "exchange_id" | "action_id">
+  input: Pick<
+    OrderWorkflow.UpdateExchangeShippingMethodWorkflowInput,
+    "exchange_id" | "action_id"
+  >
 }
 
 /**
@@ -132,8 +138,10 @@ export const updateExchangeShippingMethodWorkflowId =
 export const updateExchangeShippingMethodWorkflow = createWorkflow(
   updateExchangeShippingMethodWorkflowId,
   function (
-    input: WorkflowData<OrderWorkflow.UpdateExchangeShippingMethodWorkflowInput>
+    input: WorkflowData<
-  ): WorkflowResponse<OrderPreviewDTO> {
+      OrderWorkflow.UpdateExchangeShippingMethodWorkflowInput & AdditionalData
+    >
+  ) {
     const orderExchange: OrderExchangeDTO = useRemoteQueryStep({
       entry_point: "order_exchange",
       fields: [
@@ -161,6 +169,19 @@ export const updateExchangeShippingMethodWorkflow = createWorkflow(
       list: false,
     }).config({ name: "order-change-query" })
 
+    const setPricingContext = createHook(
+      "setPricingContext",
+      {
+        order_exchange: orderExchange,
+        order_change: orderChange,
+        additional_data: input.additional_data,
+      },
+      {
+        resultValidator: pricingContextResult,
+      }
+    )
+    const setPricingContextResult = setPricingContext.getResult()
+
     const shippingOptions = when({ input }, ({ input }) => {
       return input.data?.custom_amount === null
     }).then(() => {
@@ -178,6 +199,18 @@ export const updateExchangeShippingMethodWorkflow = createWorkflow(
         }
       )
 
+      const pricingContext = transform(
+        { action, setPricingContextResult },
+        (data) => {
+          return {
+            ...(data.setPricingContextResult
+              ? data.setPricingContextResult
+              : {}),
+            currency_code: data.action.currency_code,
+          }
+        }
+      )
+
       const shippingMethod = useRemoteQueryStep({
         entry_point: "order_shipping_method",
         fields: ["id", "shipping_option_id"],
@@ -198,7 +231,7 @@ export const updateExchangeShippingMethodWorkflow = createWorkflow(
         variables: {
           id: shippingMethod.shipping_option_id,
           calculated_price: {
-            context: { currency_code: action.currency_code },
+            context: pricingContext,
           },
         },
       }).config({ name: "fetch-shipping-option" })
@@ -220,6 +253,11 @@ export const updateExchangeShippingMethodWorkflow = createWorkflow(
       updateOrderShippingMethodsStep([updateData.shippingMethod!])
     )
 
-    return new WorkflowResponse(previewOrderChangeStep(orderExchange.order_id))
+    return new WorkflowResponse(
+      previewOrderChangeStep(orderExchange.order_id) as OrderPreviewDTO,
+      {
+        hooks: [setPricingContext] as const,
+      }
+    )
   }
 )

packages/core/core-flows/src/order/workflows/order-edit/create-order-edit-shipping-method.ts

@@ -1,4 +1,5 @@
 import {
+  AdditionalData,
   BigNumberInput,
   OrderChangeDTO,
   OrderDTO,
@@ -7,6 +8,7 @@ import {
 import { ChangeActionType, OrderChangeStatus } from "@medusajs/framework/utils"
 import {
   WorkflowResponse,
+  createHook,
   createStep,
   createWorkflow,
   transform,
@@ -21,6 +23,7 @@ import {
 import { prepareShippingMethod } from "../../utils/prepare-shipping-method"
 import { createOrderChangeActionsWorkflow } from "../create-order-change-actions"
 import { updateOrderTaxLinesWorkflow } from "../update-tax-lines"
+import { pricingContextResult } from "../../../cart/utils/schemas"
 
 /**
  * The data to validate that a shipping method can be created for an order edit.
@@ -114,8 +117,8 @@ export const createOrderEditShippingMethodWorkflowId =
 export const createOrderEditShippingMethodWorkflow = createWorkflow(
   createOrderEditShippingMethodWorkflowId,
   function (
-    input: CreateOrderEditShippingMethodWorkflowInput
+    input: CreateOrderEditShippingMethodWorkflowInput & AdditionalData
-  ): WorkflowResponse<OrderPreviewDTO> {
+  ) {
     const order: OrderDTO = useRemoteQueryStep({
       entry_point: "orders",
       fields: ["id", "status", "currency_code", "canceled_at"],
@@ -124,6 +127,29 @@ export const createOrderEditShippingMethodWorkflow = createWorkflow(
       throw_if_key_not_found: true,
     }).config({ name: "order-query" })
 
+    const setPricingContext = createHook(
+      "setPricingContext",
+      {
+        order,
+        shipping_option_id: input.shipping_option_id,
+        additional_data: input.additional_data,
+      },
+      {
+        resultValidator: pricingContextResult,
+      }
+    )
+    const setPricingContextResult = setPricingContext.getResult()
+
+    const pricingContext = transform(
+      { order, setPricingContextResult },
+      (data) => {
+        return {
+          ...(data.setPricingContextResult ? data.setPricingContextResult : {}),
+          currency_code: data.order.currency_code,
+        }
+      }
+    )
+
     const shippingOptions = useRemoteQueryStep({
       entry_point: "shipping_option",
       fields: [
@@ -135,7 +161,7 @@ export const createOrderEditShippingMethodWorkflow = createWorkflow(
       variables: {
         id: input.shipping_option_id,
         calculated_price: {
-          context: { currency_code: order.currency_code },
+          context: pricingContext,
         },
       },
     }).config({ name: "fetch-shipping-option" })
@@ -213,6 +239,11 @@ export const createOrderEditShippingMethodWorkflow = createWorkflow(
       input: [orderChangeActionInput],
     })
 
-    return new WorkflowResponse(previewOrderChangeStep(order.id))
+    return new WorkflowResponse(
+      previewOrderChangeStep(order.id) as OrderPreviewDTO,
+      {
+        hooks: [setPricingContext] as const,
+      }
+    )
   }
 )

packages/core/core-flows/src/order/workflows/order-edit/update-order-edit-shipping-method.ts

@@ -1,4 +1,5 @@
 import {
+  AdditionalData,
   OrderChangeActionDTO,
   OrderChangeDTO,
   OrderDTO,
@@ -9,6 +10,7 @@ import { ChangeActionType, OrderChangeStatus } from "@medusajs/framework/utils"
 import {
   WorkflowData,
   WorkflowResponse,
+  createHook,
   createStep,
   createWorkflow,
   parallelize,
@@ -23,6 +25,7 @@ import {
 import { previewOrderChangeStep } from "../../steps/preview-order-change"
 import { throwIfOrderChangeIsNotActive } from "../../utils/order-validation"
 import { prepareShippingMethodUpdate } from "../../utils/prepare-shipping-method"
+import { pricingContextResult } from "../../../cart/utils/schemas"
 
 /**
  * The data to validate that an order edit's shipping method can be updated.
@@ -35,7 +38,10 @@ export type UpdateOrderEditShippingMethodValidationStepInput = {
   /**
    * The details of the shipping method to be updated.
    */
-  input: Pick<OrderWorkflow.UpdateOrderEditShippingMethodWorkflowInput, "order_id" | "action_id">
+  input: Pick<
+    OrderWorkflow.UpdateOrderEditShippingMethodWorkflowInput,
+    "order_id" | "action_id"
+  >
 }
 
 /**
@@ -117,8 +123,10 @@ export const updateOrderEditShippingMethodWorkflowId =
 export const updateOrderEditShippingMethodWorkflow = createWorkflow(
   updateOrderEditShippingMethodWorkflowId,
   function (
-    input: WorkflowData<OrderWorkflow.UpdateOrderEditShippingMethodWorkflowInput>
+    input: WorkflowData<
-  ): WorkflowResponse<OrderPreviewDTO> {
+      OrderWorkflow.UpdateOrderEditShippingMethodWorkflowInput & AdditionalData
+    >
+  ) {
     const order: OrderDTO = useRemoteQueryStep({
       entry_point: "order_claim",
       fields: ["id", "currency_code"],
@@ -139,6 +147,19 @@ export const updateOrderEditShippingMethodWorkflow = createWorkflow(
       list: false,
     }).config({ name: "order-change-query" })
 
+    const setPricingContext = createHook(
+      "setPricingContext",
+      {
+        order: order,
+        order_change: orderChange,
+        additional_data: input.additional_data,
+      },
+      {
+        resultValidator: pricingContextResult,
+      }
+    )
+    const setPricingContextResult = setPricingContext.getResult()
+
     const shippingOptions = when({ input }, ({ input }) => {
       return input.data?.custom_amount === null
     }).then(() => {
@@ -156,6 +177,18 @@ export const updateOrderEditShippingMethodWorkflow = createWorkflow(
         }
       )
 
+      const pricingContext = transform(
+        { action, setPricingContextResult },
+        (data) => {
+          return {
+            ...(data.setPricingContextResult
+              ? data.setPricingContextResult
+              : {}),
+            currency_code: data.action.currency_code,
+          }
+        }
+      )
+
       const shippingMethod = useRemoteQueryStep({
         entry_point: "order_shipping_method",
         fields: ["id", "shipping_option_id"],
@@ -176,7 +209,7 @@ export const updateOrderEditShippingMethodWorkflow = createWorkflow(
         variables: {
           id: shippingMethod.shipping_option_id,
           calculated_price: {
-            context: { currency_code: action.currency_code },
+            context: pricingContext,
           },
         },
       }).config({ name: "fetch-shipping-option" })
@@ -197,6 +230,11 @@ export const updateOrderEditShippingMethodWorkflow = createWorkflow(
       updateOrderShippingMethodsStep([updateData.shippingMethod!])
     )
 
-    return new WorkflowResponse(previewOrderChangeStep(input.order_id))
+    return new WorkflowResponse(
+      previewOrderChangeStep(input.order_id) as OrderPreviewDTO,
+      {
+        hooks: [setPricingContext] as const,
+      }
+    )
   }
 )

packages/core/core-flows/src/order/workflows/return/create-complete-return.ts

@@ -7,6 +7,7 @@ import {
   OrderWorkflow,
   ShippingOptionDTO,
   WithCalculatedPrice,
+  AdditionalData,
 } from "@medusajs/framework/types"
 import {
   MathBN,
@@ -18,6 +19,7 @@ import {
 import {
   WorkflowData,
   WorkflowResponse,
+  createHook,
   createStep,
   createWorkflow,
   parallelize,
@@ -36,6 +38,7 @@ import {
   throwIfOrderIsCancelled,
 } from "../../utils/order-validation"
 import { validateReturnReasons } from "../../utils/validate-return-reason"
+import { pricingContextResult } from "../../../cart/utils/schemas"
 
 function prepareShippingMethodData({
   orderId,
@@ -186,6 +189,7 @@ function prepareFulfillmentData({
 function prepareReturnShippingOptionQueryVariables({
   order,
   input,
+  setPricingContextResult,
 }: {
   order: {
     currency_code: string
@@ -194,11 +198,13 @@ function prepareReturnShippingOptionQueryVariables({
   input: {
     return_shipping?: OrderWorkflow.CreateOrderReturnWorkflowInput["return_shipping"]
   }
+  setPricingContextResult?: any
 }) {
   const variables = {
     id: input.return_shipping?.option_id,
     calculated_price: {
       context: {
+        ...(setPricingContextResult ? setPricingContextResult : {}),
         currency_code: order.currency_code,
       },
     },
@@ -309,8 +315,10 @@ export const createAndCompleteReturnOrderWorkflowId =
 export const createAndCompleteReturnOrderWorkflow = createWorkflow(
   createAndCompleteReturnOrderWorkflowId,
   function (
-    input: WorkflowData<OrderWorkflow.CreateOrderReturnWorkflowInput>
+    input: WorkflowData<
-  ): WorkflowResponse<ReturnDTO> {
+      OrderWorkflow.CreateOrderReturnWorkflowInput & AdditionalData
+    >
+  ) {
     const order: OrderDTO = useRemoteQueryStep({
       entry_point: "orders",
       fields: [
@@ -329,8 +337,20 @@ export const createAndCompleteReturnOrderWorkflow = createWorkflow(
 
     createCompleteReturnValidationStep({ order, input })
 
+    const setPricingContext = createHook(
+      "setPricingContext",
+      {
+        order,
+        additional_data: input.additional_data,
+      },
+      {
+        resultValidator: pricingContextResult,
+      }
+    )
+    const setPricingContextResult = setPricingContext.getResult()
+
     const returnShippingOptionsVariables = transform(
-      { input, order },
+      { input, order, setPricingContextResult },
       prepareReturnShippingOptionQueryVariables
     )
 
@@ -415,6 +435,8 @@ export const createAndCompleteReturnOrderWorkflow = createWorkflow(
       }).config({ name: "emit-return-received-event" })
     )
 
-    return new WorkflowResponse(returnCreated)
+    return new WorkflowResponse(returnCreated as ReturnDTO, {
+      hooks: [setPricingContext] as const,
+    })
   }
 )

packages/core/core-flows/src/order/workflows/return/update-return-shipping-method.ts

@@ -1,4 +1,5 @@
 import {
+  AdditionalData,
   OrderChangeActionDTO,
   OrderChangeDTO,
   OrderPreviewDTO,
@@ -9,6 +10,7 @@ import { ChangeActionType, OrderChangeStatus } from "@medusajs/framework/utils"
 import {
   WorkflowData,
   WorkflowResponse,
+  createHook,
   createStep,
   createWorkflow,
   parallelize,
@@ -26,6 +28,7 @@ import {
   throwIfOrderChangeIsNotActive,
 } from "../../utils/order-validation"
 import { prepareShippingMethodUpdate } from "../../utils/prepare-shipping-method"
+import { pricingContextResult } from "../../../cart/utils/schemas"
 
 /**
  * The data to validate that a return's shipping method can be updated.
@@ -42,7 +45,10 @@ export type UpdateReturnShippingMethodValidationStepInput = {
   /**
    * The details of updating the shipping method.
    */
-  input: Pick<OrderWorkflow.UpdateReturnShippingMethodWorkflowInput, "return_id" | "action_id">
+  input: Pick<
+    OrderWorkflow.UpdateReturnShippingMethodWorkflowInput,
+    "return_id" | "action_id"
+  >
 }
 
 /**
@@ -128,8 +134,10 @@ export const updateReturnShippingMethodWorkflowId =
 export const updateReturnShippingMethodWorkflow = createWorkflow(
   updateReturnShippingMethodWorkflowId,
   function (
-    input: WorkflowData<OrderWorkflow.UpdateReturnShippingMethodWorkflowInput>
+    input: WorkflowData<
-  ): WorkflowResponse<OrderPreviewDTO> {
+      OrderWorkflow.UpdateReturnShippingMethodWorkflowInput & AdditionalData
+    >
+  ) {
     const orderReturn: ReturnDTO = useRemoteQueryStep({
       entry_point: "return",
       fields: [
@@ -157,6 +165,19 @@ export const updateReturnShippingMethodWorkflow = createWorkflow(
       list: false,
     }).config({ name: "order-change-query" })
 
+    const setPricingContext = createHook(
+      "setPricingContext",
+      {
+        order_return: orderReturn,
+        order_change: orderChange,
+        additional_data: input.additional_data,
+      },
+      {
+        resultValidator: pricingContextResult,
+      }
+    )
+    const setPricingContextResult = setPricingContext.getResult()
+
     const shippingOptions = when({ input }, ({ input }) => {
       return input.data?.custom_amount === null
     }).then(() => {
@@ -174,6 +195,18 @@ export const updateReturnShippingMethodWorkflow = createWorkflow(
         }
       )
 
+      const pricingContext = transform(
+        { action, setPricingContextResult },
+        (data) => {
+          return {
+            ...(data.setPricingContextResult
+              ? data.setPricingContextResult
+              : {}),
+            currency_code: data.action.currency_code,
+          }
+        }
+      )
+
       const shippingMethod = useRemoteQueryStep({
         entry_point: "order_shipping_method",
         fields: ["id", "shipping_option_id"],
@@ -194,7 +227,7 @@ export const updateReturnShippingMethodWorkflow = createWorkflow(
         variables: {
           id: shippingMethod.shipping_option_id,
           calculated_price: {
-            context: { currency_code: action.currency_code },
+            context: pricingContext,
           },
         },
       }).config({ name: "fetch-shipping-option" })
@@ -216,6 +249,11 @@ export const updateReturnShippingMethodWorkflow = createWorkflow(
       updateOrderShippingMethodsStep([updateData.shippingMethod!])
     )
 
-    return new WorkflowResponse(previewOrderChangeStep(orderReturn.order_id))
+    return new WorkflowResponse(
+      previewOrderChangeStep(orderReturn.order_id) as OrderPreviewDTO,
+      {
+        hooks: [setPricingContext] as const,
+      }
+    )
   }
 )

packages/core/workflows-sdk/src/utils/composer/create-hook.ts

@@ -85,6 +85,9 @@ export function createHook<Name extends string, TInvokeInput, TInvokeOutput>(
       if (options.resultValidator) {
         return options.resultValidator.parse(result)
       }
+      if (result === undefined) {
+        return new StepResponse(undefined)
+      }
       return result
     },
     () => void 0