chore(core-flows): apply taxes to items and shipping on RMA flows (#8858)

** What
Apply taxes to order line items and shipping methods

integration-tests/http/__tests__/claims/claims.spec.ts

@@ -11,6 +11,7 @@ import {
   adminHeaders,
   createAdminUser,
 } from "../../../helpers/create-admin-user"
+import { setupTaxStructure } from "../../../modules/__tests__/fixtures"
 
 jest.setTimeout(30000)
 
@@ -399,6 +400,8 @@ medusaIntegrationTestRunner({
       ).data.shipping_option
 
       item = order.items[0]
+
+      await setupTaxStructure(container.resolve(ModuleRegistrationName.TAX))
     })
 
     describe("Claims lifecycle", () => {
@@ -698,7 +701,7 @@ medusaIntegrationTestRunner({
           expect(paymentCollections[0]).toEqual(
             expect.objectContaining({
               status: "not_paid",
-              amount: 110.5,
+              amount: 109.5,
               currency_code: "usd",
             })
           )
@@ -742,7 +745,7 @@ medusaIntegrationTestRunner({
         })
 
         it("should create a payment collection successfully & mark as paid", async () => {
-          const paymentDelta = 110.5
+          const paymentDelta = 109.5
           const orderForPayment = (
             await api.get(`/admin/orders/${order.id}`, adminHeaders)
           ).data.order
@@ -1100,9 +1103,9 @@ medusaIntegrationTestRunner({
 
           expect(orderCheck.summary).toEqual(
             expect.objectContaining({
-              pending_difference: -10,
+              pending_difference: -11,
-              current_order_total: 51,
+              current_order_total: 50,
-              original_order_total: 61,
+              original_order_total: 60,
               temporary_difference: 15,
             })
           )

integration-tests/http/__tests__/exchanges/exchanges.spec.ts

@@ -9,6 +9,7 @@ import {
   adminHeaders,
   createAdminUser,
 } from "../../../helpers/create-admin-user"
+import { setupTaxStructure } from "../../../modules/__tests__/fixtures/tax"
 
 jest.setTimeout(30000)
 
@@ -379,6 +380,8 @@ medusaIntegrationTestRunner({
         },
         adminHeaders
       )
+
+      await setupTaxStructure(container.resolve(ModuleRegistrationName.TAX))
     })
 
     describe("Exchanges lifecycle", () => {

integration-tests/http/__tests__/order-edits/order-edits.spec.ts

@@ -16,6 +16,7 @@ jest.setTimeout(30000)
 medusaIntegrationTestRunner({
   testSuite: ({ dbConnection, getContainer, api }) => {
     let order
+    let taxLine
     let shippingOption
     let shippingProfile
     let fulfillmentSet
@@ -51,6 +52,30 @@ medusaIntegrationTestRunner({
         )
       ).data.customer
 
+      const taxRegion = (
+        await api.post(
+          "/admin/tax-regions",
+          {
+            country_code: "US",
+          },
+          adminHeaders
+        )
+      ).data.tax_region
+
+      taxLine = (
+        await api.post(
+          "/admin/tax-rates",
+          {
+            rate: 10,
+            code: "standard",
+            name: "Taxation is theft",
+            is_default: true,
+            tax_region_id: taxRegion.id,
+          },
+          adminHeaders
+        )
+      ).data.tax_rate
+
       const salesChannel = (
         await api.post(
           "/admin/sales-channels",
@@ -384,6 +409,12 @@ medusaIntegrationTestRunner({
         expect(result.summary.original_order_total).toEqual(60)
         expect(result.items.length).toEqual(2)
 
+        const newItem = result.items.find(
+          (i) => i.variant_id === productExtra.variants[0].id
+        )
+        expect(newItem.tax_lines[0].tax_rate_id).toEqual(taxLine.id)
+        expect(newItem.tax_lines[0].rate).toEqual(10)
+
         result = (
           await api.post(
             `/admin/order-edits/${orderId}/confirm`,
@@ -395,7 +426,7 @@ medusaIntegrationTestRunner({
         result = (await api.get(`/admin/orders/${orderId}`, adminHeaders)).data
           .order
 
-        expect(result.total).toEqual(34)
+        expect(result.total).toEqual(36.4)
         expect(result.items.length).toEqual(1)
 
         result = (

integration-tests/modules/__tests__/order/draft-order.spec.ts

@@ -226,7 +226,7 @@ medusaIntegrationTestRunner({
           shipping_methods: [
             {
               name: "test-method",
-              option_id: "test-option",
+              shipping_option_id: "test-option",
               amount: 100,
             },
           ],
@@ -268,7 +268,13 @@ medusaIntegrationTestRunner({
                     value: "3000",
                   }),
                   metadata: {},
-                  tax_lines: [],
+                  tax_lines: [
+                    expect.objectContaining({
+                      code: "US_DEF",
+                      provider_id: "system",
+                      rate: 2,
+                    }),
+                  ],
                   adjustments: [],
                   unit_price: 3000,
                   quantity: 2,
@@ -360,9 +366,15 @@ medusaIntegrationTestRunner({
                     value: "100",
                   }),
                   is_tax_inclusive: false,
-                  shipping_option_id: null,
+                  shipping_option_id: "test-option",
                   data: null,
-                  tax_lines: [],
+                  tax_lines: [
+                    expect.objectContaining({
+                      code: "US_DEF",
+                      provider_id: "system",
+                      rate: 2,
+                    }),
+                  ],
                   adjustments: [],
                   amount: 100,
                 }),

packages/core/core-flows/src/cart/steps/get-item-tax-lines.ts

@@ -17,11 +17,13 @@ export interface GetItemTaxLinesStepInput {
   items: CartLineItemDTO[]
   shipping_methods: CartShippingMethodDTO[]
   force_tax_calculation?: boolean
+  is_return?: boolean
 }
 
 function normalizeTaxModuleContext(
   cart: CartWorkflowDTO,
-  forceTaxCalculation: boolean
+  forceTaxCalculation: boolean,
+  isReturn?: boolean
 ): TaxCalculationContext | null {
   const address = cart.shipping_address
   const shouldCalculateTax = forceTaxCalculation || cart.region?.automatic_taxes
@@ -59,8 +61,7 @@ function normalizeTaxModuleContext(
       postal_code: address.postal_code,
     },
     customer,
-    // TODO: Should probably come in from order module, defaulting to false
+    is_return: isReturn ?? false,
-    is_return: false,
   }
 }
 
@@ -105,13 +106,18 @@ export const getItemTaxLinesStep = createStep(
       items,
       shipping_methods: shippingMethods,
       force_tax_calculation: forceTaxCalculation = false,
+      is_return: isReturn = false,
     } = data
 
     const taxService = container.resolve<ITaxModuleService>(
       ModuleRegistrationName.TAX
     )
 
-    const taxContext = normalizeTaxModuleContext(cart, forceTaxCalculation)
+    const taxContext = normalizeTaxModuleContext(
+      cart,
+      forceTaxCalculation,
+      isReturn
+    )
 
     if (!taxContext) {
       return new StepResponse({

packages/core/core-flows/src/common/steps/use-remote-query.ts

@@ -2,7 +2,7 @@ import {
   ContainerRegistrationKeys,
   remoteQueryObjectFromString,
 } from "@medusajs/utils"
-import { createStep, StepResponse } from "@medusajs/workflows-sdk"
+import { StepResponse, createStep } from "@medusajs/workflows-sdk"
 
 /**
  * The remote query's details.

packages/core/core-flows/src/order/steps/get-item-tax-lines.ts

@@ -17,13 +17,17 @@ export interface GetOrderItemTaxLinesStepInput {
   items: OrderLineItemDTO[]
   shipping_methods: OrderShippingMethodDTO[]
   force_tax_calculation?: boolean
+  is_return?: boolean
+  shipping_address?: OrderWorkflowDTO["shipping_address"]
 }
 
 function normalizeTaxModuleContext(
   order: OrderWorkflowDTO,
-  forceTaxCalculation: boolean
+  forceTaxCalculation: boolean,
+  isReturn?: boolean,
+  shippingAddress?: OrderWorkflowDTO["shipping_address"]
 ): TaxCalculationContext | null {
-  const address = order.shipping_address
+  const address = shippingAddress ?? order.shipping_address
   const shouldCalculateTax =
     forceTaxCalculation || order.region?.automatic_taxes
 
@@ -58,7 +62,7 @@ function normalizeTaxModuleContext(
       postal_code: address.postal_code,
     },
     customer,
-    is_return: false,
+    is_return: isReturn ?? false,
   }
 }
 
@@ -109,12 +113,19 @@ export const getOrderItemTaxLinesStep = createStep(
       items = [],
       shipping_methods: shippingMethods = [],
       force_tax_calculation: forceTaxCalculation = false,
+      is_return: isReturn = false,
+      shipping_address: shippingAddress,
     } = data
     const taxService = container.resolve<ITaxModuleService>(
       ModuleRegistrationName.TAX
     )
 
-    const taxContext = normalizeTaxModuleContext(order, forceTaxCalculation)
+    const taxContext = normalizeTaxModuleContext(
+      order,
+      forceTaxCalculation,
+      isReturn,
+      shippingAddress
+    )
 
     const stepResponseData = {
       lineItemTaxLines: [] as ItemTaxLineDTO[],

packages/core/core-flows/src/order/utils/prepare-custom-line-item-data.ts

@@ -1,10 +1,20 @@
-import { BigNumberInput } from "@medusajs/types"
+import {
+  BigNumberInput,
+  CreateOrderAdjustmentDTO,
+  CreateOrderLineItemTaxLineDTO,
+} from "@medusajs/types"
+import {
+  prepareAdjustmentsData,
+  prepareTaxLinesData,
+} from "../../cart/utils/prepare-line-item-data"
 
 interface Input {
   quantity: BigNumberInput
   metadata?: Record<string, any>
   unitPrice: BigNumberInput
   isTaxInclusive?: boolean
+  taxLines?: CreateOrderLineItemTaxLineDTO[]
+  adjustments?: CreateOrderAdjustmentDTO[]
   variant: {
     title: string
     sku?: string
@@ -24,7 +34,15 @@ interface Output {
 }
 
 export function prepareCustomLineItemData(data: Input): Output {
-  const { variant, unitPrice, isTaxInclusive, quantity, metadata } = data
+  const {
+    variant,
+    unitPrice,
+    isTaxInclusive,
+    quantity,
+    metadata,
+    taxLines,
+    adjustments,
+  } = data
 
   const lineItem: any = {
     quantity,
@@ -38,5 +56,13 @@ export function prepareCustomLineItemData(data: Input): Output {
     metadata,
   }
 
+  if (taxLines) {
+    lineItem.tax_lines = prepareTaxLinesData(taxLines)
+  }
+
+  if (adjustments) {
+    lineItem.adjustments = prepareAdjustmentsData(adjustments)
+  }
+
   return lineItem
 }

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

@@ -7,7 +7,6 @@ import {
   parallelize,
   transform,
 } from "@medusajs/workflows-sdk"
-import { useRemoteQueryStep } from "../../common"
 import { findOneOrAnyRegionStep } from "../../cart/steps/find-one-or-any-region"
 import { findOrCreateCustomerStep } from "../../cart/steps/find-or-create-customer"
 import { findSalesChannelStep } from "../../cart/steps/find-sales-channel"
@@ -15,6 +14,7 @@ import { getVariantPriceSetsStep } from "../../cart/steps/get-variant-price-sets
 import { validateVariantPricesStep } from "../../cart/steps/validate-variant-prices"
 import { prepareLineItemData } from "../../cart/utils/prepare-line-item-data"
 import { confirmVariantInventoryWorkflow } from "../../cart/workflows/confirm-variant-inventory"
+import { useRemoteQueryStep } from "../../common"
 import { createOrderLineItemsStep } from "../steps"
 import { productVariantsFields } from "../utils/fields"
 import { prepareCustomLineItemData } from "../utils/prepare-custom-line-item-data"
@@ -29,9 +29,13 @@ function prepareLineItems(data) {
           ...item,
         },
         unitPrice: MathBN.max(0, item.unit_price),
-        isTaxInclusive: item.is_tax_inclusive,
+        isTaxInclusive:
+          item.is_tax_inclusive ??
+          data.priceSets[item.variant_id!]?.is_calculated_price_tax_inclusive,
         quantity: item.quantity as number,
-        metadata: item?.metadata ?? {},
+        metadata: item?.metadata,
+        taxLines: item.tax_lines || [],
+        adjustments: item.adjustments || [],
       })
     }
 
@@ -46,7 +50,7 @@ function prepareLineItems(data) {
         item.is_tax_inclusive ??
         data.priceSets[item.variant_id!]?.is_calculated_price_tax_inclusive,
       quantity: item.quantity as number,
-      metadata: item?.metadata ?? {},
+      metadata: item?.metadata,
       taxLines: item.tax_lines || [],
       adjustments: item.adjustments || [],
     })

packages/core/core-flows/src/order/workflows/claim/claim-add-new-item.ts

@@ -21,6 +21,7 @@ import {
 } from "../../utils/order-validation"
 import { addOrderLineItemsWorkflow } from "../add-line-items"
 import { createOrderChangeActionsWorkflow } from "../create-order-change-actions"
+import { updateOrderTaxLinesWorkflow } from "../update-tax-lines"
 
 /**
  * This step validates that a new item can be added to the claim.
@@ -93,6 +94,17 @@ export const orderClaimAddNewItemWorkflow = createWorkflow(
       },
     })
 
+    const lineItemIds = transform(lineItems, (lineItems) => {
+      return lineItems.map((item) => item.id)
+    })
+
+    updateOrderTaxLinesWorkflow.runAsStep({
+      input: {
+        order_id: order.id,
+        item_ids: lineItemIds,
+      },
+    })
+
     const orderChangeActionInput = transform(
       { order, orderChange, orderClaim, items: input.items, lineItems },
       ({ order, orderChange, orderClaim, items, lineItems }) => {

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

@@ -20,6 +20,7 @@ import {
   throwIfOrderChangeIsNotActive,
 } from "../../utils/order-validation"
 import { createOrderChangeActionsWorkflow } from "../create-order-change-actions"
+import { updateOrderTaxLinesWorkflow } from "../update-tax-lines"
 
 /**
  * This step confirms that a shipping method can be created for a claim.
@@ -132,6 +133,22 @@ export const createClaimShippingMethodWorkflow = createWorkflow(
       shipping_methods: [shippingMethodInput],
     })
 
+    const shippingMethodIds = transform(createdMethods, (createdMethods) => {
+      return createdMethods.map((item) => item.id)
+    })
+
+    const isReturn = transform(input, (data) => {
+      return !!data.return_id
+    })
+
+    updateOrderTaxLinesWorkflow.runAsStep({
+      input: {
+        order_id: order.id,
+        shipping_method_ids: shippingMethodIds,
+        is_return: isReturn,
+      },
+    })
+
     const orderChangeActionInput = transform(
       {
         order,

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

@@ -20,6 +20,7 @@ import {
   throwIfOrderChangeIsNotActive,
 } from "../../utils/order-validation"
 import { createOrderChangeActionsWorkflow } from "../create-order-change-actions"
+import { updateOrderTaxLinesWorkflow } from "../update-tax-lines"
 
 /**
  * This step validates that a shipping method can be created for an exchange.
@@ -136,6 +137,22 @@ export const createExchangeShippingMethodWorkflow = createWorkflow(
       shipping_methods: [shippingMethodInput],
     })
 
+    const shippingMethodIds = transform(createdMethods, (createdMethods) => {
+      return createdMethods.map((item) => item.id)
+    })
+
+    const isReturn = transform(input, (data) => {
+      return !!data.return_id
+    })
+
+    updateOrderTaxLinesWorkflow.runAsStep({
+      input: {
+        order_id: order.id,
+        shipping_method_ids: shippingMethodIds,
+        is_return: isReturn,
+      },
+    })
+
     const orderChangeActionInput = transform(
       {
         order,

packages/core/core-flows/src/order/workflows/exchange/exchange-add-new-item.ts

@@ -21,6 +21,7 @@ import {
 } from "../../utils/order-validation"
 import { addOrderLineItemsWorkflow } from "../add-line-items"
 import { createOrderChangeActionsWorkflow } from "../create-order-change-actions"
+import { updateOrderTaxLinesWorkflow } from "../update-tax-lines"
 
 /**
  * This step validates that new items can be added to an exchange.
@@ -93,6 +94,17 @@ export const orderExchangeAddNewItemWorkflow = createWorkflow(
       },
     })
 
+    const lineItemIds = transform(lineItems, (lineItems) => {
+      return lineItems.map((item) => item.id)
+    })
+
+    updateOrderTaxLinesWorkflow.runAsStep({
+      input: {
+        order_id: order.id,
+        item_ids: lineItemIds,
+      },
+    })
+
     const orderChangeActionInput = transform(
       { order, orderChange, orderExchange, items: input.items, lineItems },
       ({ order, orderChange, orderExchange, items, lineItems }) => {

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

@@ -19,6 +19,7 @@ import {
   throwIfOrderChangeIsNotActive,
 } from "../../utils/order-validation"
 import { createOrderChangeActionsWorkflow } from "../create-order-change-actions"
+import { updateOrderTaxLinesWorkflow } from "../update-tax-lines"
 
 /**
  * This step validates that a shipping method can be created for an order edit.
@@ -113,6 +114,17 @@ export const createOrderEditShippingMethodWorkflow = createWorkflow(
       shipping_methods: [shippingMethodInput],
     })
 
+    const shippingMethodIds = transform(createdMethods, (createdMethods) => {
+      return createdMethods.map((item) => item.id)
+    })
+
+    updateOrderTaxLinesWorkflow.runAsStep({
+      input: {
+        order_id: order.id,
+        shipping_method_ids: shippingMethodIds,
+      },
+    })
+
     const orderChangeActionInput = transform(
       {
         order,

packages/core/core-flows/src/order/workflows/order-edit/order-edit-add-new-item.ts

@@ -20,6 +20,7 @@ import {
 } from "../../utils/order-validation"
 import { addOrderLineItemsWorkflow } from "../add-line-items"
 import { createOrderChangeActionsWorkflow } from "../create-order-change-actions"
+import { updateOrderTaxLinesWorkflow } from "../update-tax-lines"
 
 /**
  * This step validates that new items can be added to an order edit.
@@ -79,6 +80,17 @@ export const orderEditAddNewItemWorkflow = createWorkflow(
       },
     })
 
+    const lineItemIds = transform(lineItems, (lineItems) => {
+      return lineItems.map((item) => item.id)
+    })
+
+    updateOrderTaxLinesWorkflow.runAsStep({
+      input: {
+        order_id: order.id,
+        item_ids: lineItemIds,
+      },
+    })
+
     const orderChangeActionInput = transform(
       { order, orderChange, items: input.items, lineItems },
       ({ order, orderChange, items, lineItems }) => {

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

@@ -20,6 +20,7 @@ import {
   throwIfOrderChangeIsNotActive,
 } from "../../utils/order-validation"
 import { createOrderChangeActionsWorkflow } from "../create-order-change-actions"
+import { updateOrderTaxLinesWorkflow } from "../update-tax-lines"
 
 /**
  * This step validates that a shipping method can be created for a return.
@@ -138,6 +139,18 @@ export const createReturnShippingMethodWorkflow = createWorkflow(
       shipping_methods: [shippingMethodInput],
     })
 
+    const shippingMethodIds = transform(createdMethods, (createdMethods) => {
+      return createdMethods.map((item) => item.id)
+    })
+
+    updateOrderTaxLinesWorkflow.runAsStep({
+      input: {
+        order_id: order.id,
+        shipping_method_ids: shippingMethodIds,
+        is_return: true,
+      },
+    })
+
     const orderChangeActionInput = transform(
       {
         order,

packages/core/core-flows/src/order/workflows/update-tax-lines.ts

@@ -1,8 +1,9 @@
-import { OrderLineItemDTO, OrderShippingMethodDTO } from "@medusajs/types"
+import { OrderWorkflowDTO } from "@medusajs/types"
 import {
   WorkflowData,
   createWorkflow,
   transform,
+  when,
 } from "@medusajs/workflows-sdk"
 import { useRemoteQueryStep } from "../../common"
 import {
@@ -10,13 +11,14 @@ import {
   setOrderTaxLinesForItemsStep,
 } from "../steps"
 
-const orderFields = [
+const completeOrderFields = [
   "id",
   "currency_code",
   "email",
   "region.id",
   "region.automatic_taxes",
   "items.id",
+  "items.is_tax_inclusive",
   "items.variant_id",
   "items.product_id",
   "items.product_title",
@@ -36,6 +38,34 @@ const orderFields = [
   "items.tax_lines.code",
   "items.tax_lines.rate",
   "items.tax_lines.provider_id",
+  "shipping_methods.id",
+  "shipping_methods.is_tax_inclusive",
+  "shipping_methods.shipping_option_id",
+  "shipping_methods.amount",
+  "shipping_methods.tax_lines.id",
+  "shipping_methods.tax_lines.description",
+  "shipping_methods.tax_lines.code",
+  "shipping_methods.tax_lines.rate",
+  "shipping_methods.tax_lines.provider_id",
+  "customer.id",
+  "customer.email",
+  "customer.groups.id",
+  "shipping_address.id",
+  "shipping_address.address_1",
+  "shipping_address.address_2",
+  "shipping_address.city",
+  "shipping_address.postal_code",
+  "shipping_address.country_code",
+  "shipping_address.region_code",
+  "shipping_address.province",
+]
+
+const orderFields = [
+  "id",
+  "currency_code",
+  "email",
+  "region.id",
+  "region.automatic_taxes",
   "shipping_methods.tax_lines.id",
   "shipping_methods.tax_lines.description",
   "shipping_methods.tax_lines.code",
@@ -56,11 +86,48 @@ const orderFields = [
   "shipping_address.province",
 ]
 
+const shippingMethodFields = [
+  "id",
+  "shipping_option_id",
+  "is_tax_inclusive",
+  "amount",
+  "tax_lines.id",
+  "tax_lines.description",
+  "tax_lines.code",
+  "tax_lines.rate",
+  "tax_lines.provider_id",
+]
+
+const lineItemFields = [
+  "id",
+  "variant_id",
+  "product_id",
+  "is_tax_inclusive",
+  "product_title",
+  "product_description",
+  "product_subtitle",
+  "product_type",
+  "product_collection",
+  "product_handle",
+  "variant_sku",
+  "variant_barcode",
+  "variant_title",
+  "title",
+  "quantity",
+  "unit_price",
+  "tax_lines.id",
+  "tax_lines.description",
+  "tax_lines.code",
+  "tax_lines.rate",
+  "tax_lines.provider_id",
+]
 export type UpdateOrderTaxLinesWorkflowInput = {
   order_id: string
-  items?: OrderLineItemDTO[]
+  item_ids?: string[]
-  shipping_methods?: OrderShippingMethodDTO[]
+  shipping_method_ids?: string[]
   force_tax_calculation?: boolean
+  is_return?: boolean
+  shipping_address?: OrderWorkflowDTO["shipping_address"]
 }
 
 export const updateOrderTaxLinesWorkflowId = "update-order-tax-lines"
@@ -69,21 +136,66 @@ export const updateOrderTaxLinesWorkflowId = "update-order-tax-lines"
  */
 export const updateOrderTaxLinesWorkflow = createWorkflow(
   updateOrderTaxLinesWorkflowId,
-  (input: WorkflowData<UpdateOrderTaxLinesWorkflowInput>): WorkflowData<void> => {
+  (
+    input: WorkflowData<UpdateOrderTaxLinesWorkflowInput>
+  ): WorkflowData<void> => {
+    const isFullOrder = transform(input, (data) => {
+      return !data.item_ids && !data.shipping_method_ids
+    })
+
+    const fetchOrderFields = transform(isFullOrder, (isFullOrder) => {
+      return isFullOrder ? completeOrderFields : orderFields
+    })
+
     const order = useRemoteQueryStep({
       entry_point: "order",
-      fields: orderFields,
+      fields: fetchOrderFields,
       variables: { id: input.order_id },
+      list: false,
+    })
+
+    const items = when({ input }, ({ input }) => {
+      return input.item_ids!?.length > 0
+    }).then(() => {
+      return useRemoteQueryStep({
+        entry_point: "order_line_item",
+        fields: lineItemFields,
+        variables: { id: input.item_ids },
+      }).config({ name: "query-order-line-items" })
+    })
+
+    const shippingMethods = when({ input }, ({ input }) => {
+      return input.shipping_method_ids!?.length > 0
+    }).then(() => {
+      return useRemoteQueryStep({
+        entry_point: "order_shipping_method",
+        fields: shippingMethodFields,
+        variables: { id: input.shipping_method_ids },
+      }).config({ name: "query-order-shipping-methods" })
     })
 
     const taxLineItems = getOrderItemTaxLinesStep(
-      transform({ input, order }, (data) => ({
+      transform(
+        { input, order, items, shippingMethods, isFullOrder },
+        (data) => {
+          const shippingMethods = data.isFullOrder
+            ? data.order.shipping_methods
+            : data.shippingMethods ?? []
+
+          const lineItems = data.isFullOrder
+            ? data.order.items
+            : data.items ?? []
+
+          return {
             order: data.order,
-        items: data.input.items || data.order.items,
+            items: lineItems,
-        shipping_methods:
+            shipping_methods: shippingMethods,
-          data.input.shipping_methods || data.order.shipping_methods,
             force_tax_calculation: data.input.force_tax_calculation,
-      }))
+            is_return: data.input.is_return ?? false,
+            shipping_address: data.input.shipping_address,
+          }
+        }
+      )
     )
 
     setOrderTaxLinesForItemsStep({

packages/medusa/src/api/admin/draft-orders/validators.ts

@@ -1,9 +1,9 @@
 import { z } from "zod"
 import { AddressPayload, BigNumberInput } from "../../utils/common-validators"
 import {
+  WithAdditionalData,
   createFindParams,
   createSelectParams,
-  WithAdditionalData,
 } from "../../utils/validators"
 
 export type AdminGetOrderParamsType = z.infer<typeof AdminGetOrderParams>
@@ -30,7 +30,7 @@ const ShippingMethod = z.object({
   shipping_method_id: z.string().nullish(),
   order_id: z.string().nullish(),
   name: z.string(),
-  option_id: z.string(),
+  shipping_option_id: z.string(),
   data: z.record(z.string(), z.unknown()).optional(),
   amount: BigNumberInput,
 })

packages/modules/tax/src/services/tax-module-service.ts

@@ -399,10 +399,7 @@ export default class TaxModuleService
 
     const parentRegion = regions.find((r) => r.province_code === null)
     if (!parentRegion) {
-      throw new MedusaError(
+      return []
-        MedusaError.Types.INVALID_DATA,
-        "No parent region found for country"
-      )
     }
 
     const toReturn = await promiseAll(