From dec43eac8a472c9cef95ca4795efb4d198b9b87d Mon Sep 17 00:00:00 2001
From: appunni-m <64424476+appunni-m@users.noreply.github.com>
Date: Mon, 4 Dec 2023 15:23:54 +0530
Subject: [PATCH] Check for __$i in the reference name

---
 .../conductor/common/utils/TaskUtils.java     | 23 +++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/common/src/main/java/com/netflix/conductor/common/utils/TaskUtils.java b/common/src/main/java/com/netflix/conductor/common/utils/TaskUtils.java
index a0521b19f1..d91f77399d 100644
--- a/common/src/main/java/com/netflix/conductor/common/utils/TaskUtils.java
+++ b/common/src/main/java/com/netflix/conductor/common/utils/TaskUtils.java
@@ -27,9 +27,24 @@ public static String getLoopOverTaskRefNameSuffix(int iteration) {
     public static String removeIterationFromTaskRefName(String referenceTaskName) {
         String[] tokens = referenceTaskName.split(TaskUtils.LOOP_TASK_DELIMITER);
         int length = tokens.length;
-        return length > 1 ? String.join(
-                TaskUtils.LOOP_TASK_DELIMITER,
-                Arrays.copyOf(tokens, length - 1)
-        ) : referenceTaskName;
+
+        // Check if the last element is an integer
+        if (length > 1 && isInteger(tokens[length - 1])) {
+            // Join all elements except the last one
+            return String.join(TaskUtils.LOOP_TASK_DELIMITER, Arrays.copyOf(tokens, length - 1));
+        } else {
+            // No integer at the end, return the original string
+            return referenceTaskName;
+        }
+    }
+
+    // Helper method to check if a string is an integer
+    private static boolean isInteger(String s) {
+        try {
+            Integer.parseInt(s);
+            return true;
+        } catch (NumberFormatException e) {
+            return false;
+        }
     }
 }