feat: Add types to the workflow conditional (#7979)

**What**
Add typing to the when conditional workflow step

FIXES CORE-2494

packages/core/workflows-sdk/package.json

@@ -39,6 +39,6 @@
     "build": "rimraf dist && tsc --build",
     "watch": "tsc --build --watch",
     "test": "jest --runInBand --bail --forceExit",
-    "test:run": "../../node_modules/.bin/ts-node ./src/utils/_playground.ts"
+    "test:run": "../../../node_modules/.bin/ts-node ./src/utils/_playground.ts"
   }
 }

packages/core/workflows-sdk/src/utils/_playground.ts

@@ -2,6 +2,8 @@ import {
   createStep,
   createWorkflow,
   StepResponse,
+  transform,
+  when,
   WorkflowData,
 } from "./composer"
 
@@ -29,7 +31,16 @@ const workflow = createWorkflow(
 
 const workflow2 = createWorkflow("workflow", function () {
   const step1Res = step1()
-  step3()
+
+  const step3Res = when({ value: true }, ({ value }) => {
+    return value
+  }).then(() => {
+    return step3()
+  })
+
+  transform({ step3Res }, ({ step3Res }) => {
+    console.log(step3Res)
+  })
 
   const workflowRes = workflow.asStep({ outsideWorkflowData: step1Res.step1 })
 

packages/core/workflows-sdk/src/utils/composer/__tests__/index.spec.ts

@@ -41,16 +41,19 @@ describe("Workflow composer", () => {
       expect(result).toEqual({ result: "hi from outside" })
     })
 
-    it("should skip step if condition is true", async function () {
+    it("should skip step if condition is false", async function () {
       const step1 = createStep("step1", async (_, context) => {
         return new StepResponse({ result: "step1" })
       })
       const step2 = createStep("step2", async (input: string, context) => {
         return new StepResponse({ result: input })
       })
-      const step3 = createStep("step3", async (input: string, context) => {
+      const step3 = createStep(
-        return new StepResponse({ result: input ?? "default response" })
+        "step3",
-      })
+        async (input: string | undefined, context) => {
+          return new StepResponse({ result: input ?? "default response" })
+        }
+      )
 
       const subWorkflow = createWorkflow(
         getNewWorkflowId(),
@@ -80,6 +83,48 @@ describe("Workflow composer", () => {
       expect(result).toEqual({ result: "default response" })
     })
 
+    it("should not skip step if condition is true", async function () {
+      const step1 = createStep("step1", async (_, context) => {
+        return new StepResponse({ result: "step1" })
+      })
+      const step2 = createStep("step2", async (input: string, context) => {
+        return new StepResponse({ result: input })
+      })
+      const step3 = createStep(
+        "step3",
+        async (input: string | undefined, context) => {
+          return new StepResponse({ result: input ?? "default response" })
+        }
+      )
+
+      const subWorkflow = createWorkflow(
+        getNewWorkflowId(),
+        function (input: WorkflowData<string>) {
+          step1()
+          return step2(input)
+        }
+      )
+
+      const workflow = createWorkflow(
+        getNewWorkflowId(),
+        function (input: { callSubFlow: boolean }) {
+          const subWorkflowRes = when({ input }, ({ input }) => {
+            return input.callSubFlow
+          }).then(() => {
+            return subWorkflow.runAsStep({
+              input: "hi from outside",
+            })
+          })
+
+          return step3(subWorkflowRes.result)
+        }
+      )
+
+      const { result } = await workflow.run({ input: { callSubFlow: true } })
+
+      expect(result).toEqual({ result: "hi from outside" })
+    })
+
     it("should revert the workflow and sub workflow on failure", async function () {
       const step1Mock = jest.fn()
       const step1 = createStep(

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

@@ -4,9 +4,9 @@ import {
   WorkflowStepHandler,
   WorkflowStepHandlerArguments,
 } from "@medusajs/orchestration"
-import { OrchestrationUtils, deepCopy, isString } from "@medusajs/utils"
+import { deepCopy, isString, OrchestrationUtils } from "@medusajs/utils"
 import { ulid } from "ulid"
-import { StepResponse, resolveValue } from "./helpers"
+import { resolveValue, StepResponse } from "./helpers"
 import { proxify } from "./helpers/proxy"
 import {
   CreateWorkflowComposerContext,
@@ -346,7 +346,7 @@ function wrapConditionalStep(
   const originalInvoke = handle.invoke
   handle.invoke = async (stepArguments: WorkflowStepHandlerArguments) => {
     const args = await resolveValue(input, stepArguments)
-    const canContinue = await condition(args)
+    const canContinue = await condition(args, stepArguments)
 
     if (stepArguments.step.definition?.async) {
       stepArguments.step.definition.backgroundExecution = true

packages/core/workflows-sdk/src/utils/composer/when.ts

@@ -1,4 +1,27 @@
 import { OrchestrationUtils } from "@medusajs/utils"
+import { StepExecutionContext, WorkflowData } from "./type"
+
+type ConditionFunction<T extends object | WorkflowData> = (
+  input: T extends WorkflowData<infer U>
+    ? U
+    : T extends object
+    ? { [K in keyof T]: T[K] extends WorkflowData<infer U> ? U : T[K] }
+    : {},
+  context: StepExecutionContext
+) => boolean
+
+type ThenFunc = <ThenResolver extends () => any>(
+  resolver: ThenResolver
+) => ReturnType<ThenResolver> extends WorkflowData<infer ReturnedWorkflowData>
+  ? Partial<WorkflowData<ReturnedWorkflowData>>
+  : ReturnType<ThenResolver>
+
+export function when<T extends object | WorkflowData, Then extends Function>(
+  values: T,
+  condition: ConditionFunction<T>
+): {
+  then: ThenFunc
+}
 
 export function when(input, condition) {
   global[OrchestrationUtils.SymbolMedusaWorkflowComposerCondition] = {