From 010560fd2a4c12eb7a551919df33e58e7cbe920f Mon Sep 17 00:00:00 2001
From: Adrien de Peretti <adrien.deperetti@gmail.com>
Date: Thu, 23 Nov 2023 10:11:03 +0100
Subject: [PATCH] fix(workflows): compensation handling (#5691)

---
 .changeset/small-apes-own.md                  |  5 +++
 .../workflows/utils/composer/compose.ts       | 37 +++++++++++++++++++
 .../src/utils/composer/create-step.ts         |  4 +-
 3 files changed, 44 insertions(+), 2 deletions(-)
 create mode 100644 .changeset/small-apes-own.md

diff --git a/.changeset/small-apes-own.md b/.changeset/small-apes-own.md
new file mode 100644
index 0000000000..38fbbbbe16
--- /dev/null
+++ b/.changeset/small-apes-own.md
@@ -0,0 +1,5 @@
+---
+"@medusajs/workflows": patch
+---
+
+Fix(workflows): compensation handling
diff --git a/integration-tests/plugins/__tests__/workflows/utils/composer/compose.ts b/integration-tests/plugins/__tests__/workflows/utils/composer/compose.ts
index c1173e8e86..0bec8edef9 100644
--- a/integration-tests/plugins/__tests__/workflows/utils/composer/compose.ts
+++ b/integration-tests/plugins/__tests__/workflows/utils/composer/compose.ts
@@ -1842,4 +1842,41 @@ describe("Workflow composer", function () {
       })
     })
   })
+
+  it("should compose a workflow that throws without crashing and the compensation will receive undefined for the step that fails", async () => {
+    const mockStep1Fn = jest.fn().mockImplementation(function (input) {
+      throw new Error("invoke fail")
+    })
+
+    const mockCompensateSte1 = jest.fn().mockImplementation(function (input) {
+      return input
+    })
+
+    const step1 = createStep("step1", mockStep1Fn, mockCompensateSte1)
+
+    const workflow = createWorkflow("workflow1", function (input) {
+      return step1(input)
+    })
+
+    const workflowInput = { test: "payload1" }
+    const { errors } = await workflow().run({
+      input: workflowInput,
+      throwOnError: false,
+    })
+
+    expect(errors).toHaveLength(1)
+    expect(errors[0]).toEqual({
+      action: "step1",
+      handlerType: "invoke",
+      error: expect.objectContaining({
+        message: "invoke fail",
+      }),
+    })
+
+    expect(mockStep1Fn).toHaveBeenCalledTimes(1)
+    expect(mockCompensateSte1).toHaveBeenCalledTimes(1)
+
+    expect(mockStep1Fn.mock.calls[0][0]).toEqual(workflowInput)
+    expect(mockCompensateSte1.mock.calls[0][0]).toEqual(undefined)
+  })
 })
diff --git a/packages/workflows/src/utils/composer/create-step.ts b/packages/workflows/src/utils/composer/create-step.ts
index 5405559c6a..7e8b0df926 100644
--- a/packages/workflows/src/utils/composer/create-step.ts
+++ b/packages/workflows/src/utils/composer/create-step.ts
@@ -33,7 +33,7 @@ type InvokeFn<TInput extends object, TOutput, TCompensateInput> = (
     >>
 
 type CompensateFn<T> = (
-  input: T,
+  input: T | undefined,
   context: StepExecutionContext
 ) => unknown | Promise<unknown>
 
@@ -121,7 +121,7 @@ function applyStep<
               context: transactionContext.context,
             }
 
-            const stepOutput = transactionContext.invoke[stepName].output
+            const stepOutput = transactionContext.invoke[stepName]?.output
             const invokeResult =
               stepOutput?.__type === SymbolWorkflowStepResponse
                 ? stepOutput.compensateInput &&