swap handles on config

Author: carlos-r-l-rodrigues

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

@@ -1413,16 +1413,34 @@ describe("Workflow composer", function () {
         return new StepResponse(input)
       })
 
-      const fakeStepWorkflow = createWorkflow("fake-workflow", () => {
+      const sameStepWorkflow = createWorkflow("fake-workflow", () => {
         const a = log(1).config({ name: "aaaa" })
         const b = log(2) // without config on purpose
         const c = log(3).config({ name: "cccc" })
         return new WorkflowResponse([a, b, c])
       })
 
-      const { result } = await fakeStepWorkflow().run()
+      const sameStepWorkflow2 = createWorkflow("fake-workflow-2", () => {
+        const a = log(1)
+        const b = log(2).config({ name: "bbbb" })
+        const c = log(3).config({ name: "cccc" })
+        return new WorkflowResponse([a, b, c])
+      })
+
+      const sameStepWorkflow3 = createWorkflow("fake-workflow-3", () => {
+        const a = log(1).config({ name: "aaaa" })
+        const b = log(2).config({ name: "bbbb" })
+        const c = log(3)
+        return new WorkflowResponse([a, b, c])
+      })
+
+      const { result } = await sameStepWorkflow().run()
+      const { result: result2 } = await sameStepWorkflow2().run()
+      const { result: result3 } = await sameStepWorkflow3().run()
 
       expect(result).toEqual([1, 2, 3])
+      expect(result2).toEqual([1, 2, 3])
+      expect(result3).toEqual([1, 2, 3])
     })
 
     it("should skip steps until the named step in case of permanent failure", async () => {

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

@@ -143,10 +143,9 @@ export function applyStep<
 
     this.isAsync ||= !!(stepConfig.async || stepConfig.compensateAsync)
 
-    if (!this.canOverrideHandlers.has(stepName)) {
-      this.canOverrideHandlers.add(stepName)
-      this.handlers.set(stepName, handler)
-    }
+    this.overriddenHandler.set(stepName, this.handlers.get(stepName)!)
+
+    this.handlers.set(stepName, handler)
 
     const ret = {
       __type: OrchestrationUtils.SymbolWorkflowStep,
@@ -179,7 +178,6 @@ export function applyStep<
       }
 
       delete newConfig.name
-      this.canOverrideHandlers.delete(stepName)
 
       const handler = createStepHandler.bind(this)({
         stepName: newStepName,
@@ -190,6 +188,9 @@ export function applyStep<
 
       wrapAsyncHandler(newConfig, handler)
 
+      this.handlers.set(stepName, this.overriddenHandler.get(stepName)!)
+      this.overriddenHandler.delete(stepName)
+
       this.handlers.set(newStepName, handler)
 
       this.flow.replaceAction(stepConfig.uuid!, newStepName, newConfig)

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

@@ -117,7 +117,7 @@ export function createWorkflow<TData, TResult, THooks extends any[]>(
     flow: WorkflowManager.getEmptyTransactionDefinition(),
     isAsync: false,
     handlers,
-    canOverrideHandlers: new Set(),
+    overriddenHandler: new Map(),
     hooks_: {
       declared: [],
       registered: [],

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

@@ -113,7 +113,7 @@ export type CreateWorkflowComposerContext = {
   flow: OrchestratorBuilder
   isAsync: boolean
   handlers: WorkflowHandler
-  canOverrideHandlers: Set<string>
+  overriddenHandler: WorkflowHandler
   stepBinder: <TOutput = unknown>(
     fn: StepFunctionResult
   ) => WorkflowData<TOutput>